【No.30】feat(dr-grpo): add Dr.GRPO support - #239
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds Dr.GRPO support to Relax’s Megatron backend, implementing fixed-budget loss reduction (via optimizer-window metadata) and Dr.GRPO’s reward centering semantics, plus accompanying tests, docs, and an example training script.
Changes:
- Register
dr_grpoand reuse the existing GRPO service topology (Rollout/Actor/Advantages/Reference/ActorFwd). - Implement Dr.GRPO reward centering (mandatory group mean subtraction; no group-std division) and fixed-budget loss scaling via optimizer-window
__loss_scale__metadata passed through the data iterator. - Add unit tests and bilingual documentation (and a runnable script) for Dr.GRPO usage and expected behavior.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/backends/megatron/test_dr_grpo.py | New unit tests for Dr.GRPO registration, reward semantics, optimizer-window metadata, CP/DP reduction behavior, and iterator replay. |
| scripts/training/text/run-qwen3-0.6B-drgrpo.sh | New colocate training recipe script for GRPO/Dr.GRPO correctness smoke tests. |
| relax/utils/utils.py | Make group reward centering mandatory for dr_grpo without std normalization. |
| relax/utils/opd/opd_utils.py | Allow OPD loss/metrics to use an optional token reducer for per-token aggregation semantics. |
| relax/utils/arguments.py | Add dr_grpo as an --advantage-estimator choice and force Megatron per-token loss mode when selected. |
| relax/core/registry.py | Register dr_grpo to use the same role topology as GRPO-family algos. |
| relax/components/advantages.py | Add dr_grpo handling in advantage/return computation, including KL reward shaping behavior. |
| relax/backends/megatron/loss.py | Implement optimizer-window metadata provider, apply explicit loss scaling, and adjust reducers/metrics behavior for Dr.GRPO. |
| relax/backends/megatron/data.py | Add optimizer-window metadata replay on micro-batches; enforce non-empty rollout windows; persist per-window sample counts in rollout data. |
| relax/backends/megatron/actor.py | Bind optimizer-window metadata into training iterators before the backward/replay stages. |
| docs/zh/guide/dr-grpo-training.md | New Chinese Dr.GRPO training guide and implementation notes. |
| docs/en/guide/dr-grpo-training.md | New English Dr.GRPO training guide and implementation notes. |
| docs/zh/guide/configuration.md | Document dr_grpo as an --advantage-estimator option. |
| docs/en/guide/configuration.md | Document dr_grpo as an --advantage-estimator option. |
| docs/zh/examples/algorithms.md | Add Dr.GRPO algorithm reference section and links. |
| docs/en/examples/algorithms.md | Add Dr.GRPO algorithm reference section and links. |
| docs/.vitepress/config.mts | Add Dr.GRPO guide links to the docs sidebar navigation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| """Prepare optimizer-window metadata consumed by the policy loss.""" | ||
| if args.advantage_estimator == "dr_grpo": | ||
| step_stats = [] |
There was a problem hiding this comment.
已在 69e1dbd 修复。
_validate_dr_grpo_args()在启动时校验rollout_max_response_len必须是正整数;prepare_policy_optimizer_window_metadata()在执行T / (N * B)的消费边界保留同样的防御性校验;tests/utils/test_arguments_dr_grpo.py覆盖0与True,provider 测试覆盖 zero budget。
这样 CLI 与程序化/helper 调用都会 fail fast,不会再延迟到除法处报 TypeError 或 ZeroDivisionError。
| for iterator in data_iterators: | ||
| iterator.metadata_by_microbatch = metadata_by_microbatch | ||
|
|
There was a problem hiding this comment.
已在 69e1dbd 修复。
bind_optimizer_window_metadata() 现在无论绑定新 metadata 还是 clear metadata 都会把 microbatch_offset 重置为 0;get_next() 也只在实际存在 metadata 时递增该 offset。新增 rebind regression test:先消费一个 micro-batch,再绑定新 metadata,确认从新 window 的第一个 entry 开始 replay。
| self.offset = 0 | ||
| self.microbatch_offset = 0 | ||
|
|
||
| def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]: |
There was a problem hiding this comment.
已修正。当前 DataIterator.get_next() 的返回类型是 dict[str, Any],与 batch 中可能包含 list、tensor 和 optimizer-window scalar metadata 的实际 contract 一致。
| step_stats = [] | ||
| start = 0 | ||
| for step_local_sample_count in step_local_sample_counts: | ||
| end = start + step_local_sample_count | ||
| local_response_tokens = torch.stack( | ||
| [mask.sum().to(dtype=torch.float32) for mask in rollout_data["loss_masks"][start:end]] | ||
| ).sum() | ||
| step_stats.append( | ||
| torch.stack( | ||
| [ | ||
| local_response_tokens.new_tensor(step_local_sample_count), | ||
| local_response_tokens, | ||
| ] | ||
| ) | ||
| ) | ||
| start = end | ||
|
|
||
| stats = torch.stack(step_stats) | ||
| dist.all_reduce(stats, group=mpu.get_data_parallel_group(with_context_parallel=False)) | ||
| step_loss_scales = stats[:, 1] / (stats[:, 0] * args.rollout_max_response_len) |
There was a problem hiding this comment.
已在 69e1dbd 修复。
prepare_policy_optimizer_window_metadata() 现在通过 device_utils.make_current_torch_device() 选择 accelerator device,仅把每个 tensor mask 的 scalar sum() 搬到该 device,再 stack/all-reduce;不会移动完整 mask。归约继续显式使用 get_data_parallel_group(with_context_parallel=False),生成的 __dr_grpo_window_scale__ 与 CUDA loss 位于同一 device。
新增 device regression test,并保留真实 Gloo DP/CP topology test;Qwen3.5-4B 四卡 H20 端到端路径也已完成。
|
感谢这份提交,亮点:
reveiw 修改意见:
relax/backends/megatron/loss.py 的 per-token 分支:
这里没有 dr_grpo 守卫。而 loss_scale 并不是 Dr.GRPO 专有的 —— 它由 relax/utils/data/stream_dataloader.py:263/1188/1258 在 fully-async streaming 路径注入。仓库里同时使用 fully-async 和 --calculate-per-token-loss 的现有配方至少有三个: scripts/training/multimodal/run-dotsocr2-8xgpu.sh 这些任务今天在 per-token 分支是故意不乘 explicit_loss_scale 的(紧跟着的那段注释解释了原因),改动后会被静默多乘一个 scale。 一个佐证:紧邻的 metrics 缩放你是加了守卫的 ——
说明你知道这里需要守卫,只是漏在了 loss 那一行。请补上同样的条件,或者给 Dr.GRPO 的 window scale 用一个独立的 key(如 dr_grpo_window_scale),避免和 streaming 的 loss_scale 共用同一通道。
题目验收明确要求「报告 reward、length、KL 和训练稳定性」。50-step 对照表目前是 Loss / Reward / Accuracy / Grad-norm —— reward ✅、稳定性(grad norm)✅,但 length 和 KL 两列缺失。 其中 length 尤其关键:Dr.GRPO 的核心论点就是「错误回答不再无意义变长」(论文 Fig.1),这是最有说服力的一条证据,而且建议按 Correct / Incorrect 分开统计 —— 合并成一条总长度曲线看不出这个效果。你的 mixed-length oracle 测试已经在单元层证明了 reduction 行为正确,但端到端那张表还需要把长度趋势体现出来。 顺带:grad norm mean 2.94 → 0.47(GRPO → Dr.GRPO)这个 6 倍的差异建议在正文里解释一句。固定预算分母 N·B 与实际 token 数 T 的比值本来就会整体缩放梯度尺度,这不是不稳定,但读者容易误读,最好说明两臂的有效学习率是否可比。
if args.advantage_estimator == "dr_grpo" and args.kl_coef != 0: get_grpo_returns()(ppo_utils.py:410-417)只是把 reward 广播到 token 形状,不做任何 KL shaping。也就是说标准 GRPO 在 Relax 里没有 reward-side KL,而这段给 dr_grpo 单独加上了。 kl_coef != 0 时两臂就不再是「只差 Dr.GRPO 的两处修改」,对照实验的结论会被这一项污染。本次实验应该是 kl_coef=0 所以没暴露。请说明这是刻意对齐某个参考实现,还是应该移除;如果保留,请在文档和对照实验协议里显式声明。
or args.advantage_estimator == "dr_grpo" 让 Dr.GRPO 的中心化无视 --disable-rewards-normalization。语义上是对的(中心化是 Dr.GRPO 的定义),但用户传了这个 flag 却没有任何提示。建议加一条 warning,或在参数校验阶段直接拒绝这个组合。
四条都是有效的,尤其第四条(设备不匹配)建议优先处理: (1) prepare_policy_optimizer_window_metadata() 未校验 --rollout-max-response-len 为正 → T/(N·B) 会 ZeroDivision/TypeError。
|
e4b08fe to
69e1dbd
Compare
|
@li126com 感谢您的 review。已在 69e1dbd 完成代码、测试、recipe 和文档修复。Qwen3.5-4B 对照结果及曲线同步到 PR 正文; 逐项回复如下:
|
8c77097 to
69e1dbd
Compare
|
重新复核后,Dr.GRPO 的主体算法实现没有发现新的 correctness 问题:独立 estimator、reward 当前剩余问题集中在验收证据与对比报告。
Dr.GRPO 路径会将 kl_loss 乘 dr_grpo_window_scale = T/(N*B),日志汇总随后再除 T, GRPO 未应用该 scale,得到的是: 所以报告中的 0.1692 vs 0.0423 不能直接解释为 Dr.GRPO 的 policy-reference 平均 KL 更小。
正文报告了 16 条 response(1024/256/64/16)、六种 DP/CP 拓扑、两种 micro-batch size、 另外,当前两个 focused test 文件静态展开为 21 cases,与正文的 35 passed 不一致。请把
Recipe 能启动同类实验,但尚未固定模型 checkpoint revision、GSM8K split/转换方式和数据 以上不否定算法实现本身;修复重点是让回归测试和对比报告与当前代码及验收要求一致。 |
f458ed7 to
4b49a79
Compare
|
@li126com 感谢您的 review,三个问题已逐条修复。 1. KL 对比口径我们已改为按统一定义 2. 回归测试抱歉在上一次提交中忽略了相关的测试,我们现在已把 deterministic mixed-length/topology oracle 入库,合并进
3. 200-step 复现材料新增
按照 Relax 的常见做法,我们暂时没有提交直接生成表格/图片的脚本,而是在 manifest 中 |
Review result: Request changes恭喜验收通过,但仍有一个会导致 CP topology-dependent gradient 的阻塞问题,在合入前需要修改。 [P1] Fully-masked response makes the Dr.GRPO denominator depend on CP topology--rollout-sample-filter-path 是公开支持的功能。被过滤的 response 不会被移出 batch,而是将整条 loss_mask 设为零:
Dr.GRPO window metadata 使用真实的 mask.sum() 计算 T:
但传给 Megatron finalizer 的 token count 定义不一致:
见 relax/backends/megatron/cp_utils.py#L192-L225 (https://github.com/redai-infra/Relax/blob/4b49a791430d4e00ee8722874b096b6cefb25476/relax/backends/megatron/cp_utils.py#L192-L225)。 最小例子:
结果:
因此同一批数据仅改变 CP degree,梯度就会改变。在提供的实际 GPU 镜像中调用当前实现也复现了: raw_metadata_T = 2 如果整个 optimizer window 都被 mask:
建议:
现有 mixed-mask fixture 中每条 response 仍至少保留 14 个有效 token,因此没有覆盖这个情况:
|
|
@li126com 感谢您的 review,我们已经修复相关问题,并将相关的改动同步到PR正文中。 1. 改动
2. 新增单测我们新增了相关的单元测试:
这些用例在修复前的代码上于 3. 端到端验证我们重跑了之前的端到端对照(Qwen3.5-4B / 4×H20 / TP2 / CP1 / 此外,我们构造了会触发该分支的场景:训练时把部分或全部 response 的 loss_mask 置零。结果符合预期:单条 response 全 mask 时,分母精确等于真实有效 token 数;整窗全 mask 时分母为 0,跳过日志在全部 rank 触发,无除零错误。该场景需要临时改动训练代码来注入 mask,所以只作为一次性验证,测试脚本未随本 PR 提交。 |
|
感谢作者持续修改。重新按题目全部验收项审查了当前 head c3f3157。此前 whole-response mask 导致 CP1/CP2 分母不一致,以及全零窗口仍执行 optimizer/scheduler 的问题已经修复;默认论文配置下的 Dr.GRPO 数学实现目前是正确的。 当前建议仍为 Request changes,只剩下面一个实现 blocker,以及一个合入前的分支同步要求。 [P1] Dr.GRPO + --normalize-advantages 在 CP 和全零 mask 下仍不正确该选项目前在 Dr.GRPO 文档中被明确描述为支持,但实现存在三个相关问题:
建议二选一:
如果选择继续支持,建议至少补:
合入前需要同步最新 main当前 PR head 与最新 main 存在真实冲突,涉及:
最新 main 已加入 RLOO,并修改了相同的 algorithm dispatch、advantage 和 reward-processing 路径。解决冲突时需要保留 Dr.GRPO 与 RLOO 两套独立语义,不能简单选择任一侧。完成后建议重新运行 Dr.GRPO focused tests、RLOO tests 和 pre-commit run --all-files。这是分支集成要求,不是额外的 Dr.GRPO 公式问题。 其余验收项重新逐项核对后,默认 Dr.GRPO 路径已满足:
非阻塞说明对比报告 manifest 绑定的是较早的 346bb53,且只跑了 CP1,不是当前 head。它没有覆盖后来修复的 CP/full-mask 边界,但这些修改不改变正常 CP1、非空窗口的数值路径,因此现有 paired report 仍可作为题目要求的 GRPO/Dr.GRPO 对比证据。若本次 rebase 没有改变默认算法语义,不要求重新跑完整 200-step E2E;补 current-head CP2/full-mask smoke test 即可。 现有测试也没有真正调用 Megatron finalizer,并用 spy 验证 optimizer/scheduler 的 .step() 没被调用。这属于测试增强项,建议补充,但在实现已经审查清楚的前提下不单独阻塞验收。 最终结论修复或明确禁用 Dr.GRPO + --normalize-advantages,正确解决最新 main/RLOO 冲突,并通过相应回归测试后,我这边可以 Approve。除非冲突解决引入新的算法语义变化,否则没有发现其他需要阻塞合入的验收问题。 |
# 🐛 Bug Fix ## Isolate fixed-budget loss scaling - Use a dedicated Dr.GRPO optimizer-window metadata key without changing fully-async streaming loss scales - Move response-token statistics to the active device before distributed reduction - Reset optimizer metadata offsets whenever metadata is rebound or cleared ## Keep the comparison semantics focused - Remove Dr.GRPO-only reward-side KL shaping from both advantage execution paths - Validate response budgets, reward centering, KL configuration, and supported closed-window modes --- # ✅ Tests ## Cover reviewer regressions - Verify scale-key isolation, DP/CP metadata reduction, device placement, and metadata replay - Cover both reward-side KL paths and the supported Dr.GRPO argument combinations --- # 📝 Documentation ## Align guides and the paired recipe - Document fixed-window metadata, pure fully-async rejection, and explicit KL loss semantics in English and Chinese - Add a Qwen3.5-4B four-GPU paired recipe with a 200-step default
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
# 🐛 Bug Fix
## Reject `--normalize-advantages` for Dr.GRPO
- Add a guard in `_validate_dr_grpo_args` so `--advantage-estimator dr_grpo`
combined with `--normalize-advantages` fails at startup with a clear error.
- Removing the advantage variance normalization is the core of Dr.GRPO, while
the flag re-applies a global masked whitening step on top of the group-centered
advantages. Honoring both at once contradicts the estimator's own semantics.
---
# 📝 Documentation
## Align the bilingual docs with the new guard
- `docs/{en,zh}/guide/dr-grpo-training.md`: rewrite the advantage-normalization
warning block, the parameter table row, and the best-practice item to state
that the flag is rejected instead of optionally honored.
- `docs/{en,zh}/examples/algorithms.md`: update the Dr.GRPO parameter table row
to the same wording.
---
# ✅ Tests
## Cover the new validation path
- `tests/utils/test_arguments_dr_grpo.py`: add `normalize_advantages` to the
argument fixture defaults and extend `test_dr_grpo_rejects_incompatible_semantics`
with the rejected combination.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
# ✅ Tests ## Drive the real `train_one_step` with optimizer/scheduler spies - Add `test_train_one_step_skips_optimizer_and_scheduler_for_empty_window`, which runs the production `train_one_step` under Gloo with spies standing in for Megatron's optimizer and `OptimizerParamScheduler`. - The empty-window flag travels through the real `forward_step`, so the test covers the full path from micro-batch metadata to the skip decision rather than asserting on the metadata alone. - Both parameterizations are checked: a fully masked window leaves `.step()` uncalled on both objects and reports `grad_norm == 0.0` instead of `NaN`, while a normal window calls each exactly once. ## Call Megatron's real gradient finalizer - Add `test_megatron_finalizer_reaches_dr_grpo_fixed_denominator`, which feeds the window's pre-normalization gradient into `megatron.core.distributed.finalize_model_grads.finalize_model_grads` and asserts Megatron's own `1/T` scaling lands on `sum(loss) / (N * B)`. - Add `test_megatron_finalizer_leaves_gradients_alone_for_empty_window`, covering `T == 0`: Megatron guards the division itself, so the gradient comes back unscaled and finite instead of raising or producing `inf`/`NaN`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
c3f3157 to
6ac0897
Compare
|
@li126com 感谢您的 review。两个 blocker 已处理,测试增强项已补上,相关改动同步至PR正文。 1. Dr.GRPO +
|
| 配置 | 注入 | Steps | 结果 |
|---|---|---|---|
| CP2 | 无 | 10/10 | job succeeded;train/loss、train/grad_norm、train/ppo_kl、rollout/raw_reward、rollout/response_lengths 全部 finite |
| full-mask | all |
10/10 | job succeeded;每步都进入空窗口分支 |
| partial-mask | even |
10/10 | job succeeded;未触发空窗口,optimizer 正常执行 |
full-mask 组的直接证据:
relax.backends.megatron.model:1198 | WARNING
Skipping optimizer step: optimizer window has no loss-contributing tokens (rollout_id=N, step_id=0)
10 个 step 全部触发,train/grad_norm 全为 0.0 而非 NaN;ZeroDivisionError 与 Traceback 计数为 0;partial-mask 中窗口内只有一半 response 被置零时仍有有效 token,跳过分支不应触发。实测跳过警告为 0 次,optimizer 正常执行。 三组均未出现 NaN/Inf、OOM、hang 或 Ray job failure。
feat(dr-grpo): add Dr.GRPO support
What
本 PR 为 Relax 的 Megatron 训练后端增加 Dr.GRPO 支持。
主要改动包括:
--advantage-estimator中加入dr_grpo,并复用 GRPO 现有的 Rollout、Actor、Advantages、Reference 和 ActorFwd 服务拓扑。其中$N$ 是 optimizer window 内的全局 response 数,$B$ 是
rollout_max_response_len。实现上继续使用 Megatron 的 per-token loss 和 CP reduction 路径,在每个 optimizer window 上应用$T/(N B)$ 的 loss scale,将 Megatron 的 token-mean 结果转换为 Dr.GRPO 的 fixed-budget 结果。Policy gradient、entropy 和 explicit KL loss 使用相同的 reduction 语义。
cp_size == 1与cp_size > 1的有效 token 定义,并处理全局有效 token 数为 0 的 optimizer window。Why
GRPO 使用实际 response token 数归一化 loss,并使用组内 reward 标准差归一化 advantage,因此训练权重会同时受 response 长度和组内 reward 分布影响。Dr.GRPO 只对组内 reward 做中心化,并使用固定的$N B$ 归一化 Actor loss,使 loss 的归一化不再随实际 response token 总数 $T$ 变化。
Megatron 在 CP 大于 1 时要求使用 per-token loss,并在 gradient finalize 阶段统一除以全局有效 token 数。因此不能简单切换到非 per-token loss 路径,否则无法正确支持 CP。
本 PR 在保留 Megatron 原有 token normalization 和 CP reduction 的基础上,通过 optimizer-window scale 得到 Dr.GRPO 的固定预算目标,避免在 Megatron 训练引擎中加入算法特定的训练流程。
How
Reward 和 advantage
对于同一个 prompt 采样出的$G$ 条 response,Dr.GRPO 使用:
与 GRPO 不同,Dr.GRPO 不再使用组内 reward 标准差除以$A_i$ 。
去掉方差归一化是 Dr.GRPO 的核心,因此
--normalize-advantages会被参数校验拒绝:该选项会在组内中心化之后再叠加一层全局 whitening,与之矛盾。Fixed-budget loss reduction
这里的 optimizer window 是一次$N$ 和 $T$ 都在这个范围内统计。
optimizer.step()对应的完整梯度累积范围,包含该次参数更新前参与 forward/backward 的所有 micro-batch。下文的令:
rollout_max_response_len。Relax 在 optimizer window 开始前计算:
随后将组合后的 Actor loss 乘以$\alpha$ ,并复用 Megatron 原有的全局 token normalization:
这样既能满足 Dr.GRPO 的固定预算目标,也能保留 Megatron 在 CP 场景下要求的
calculate_per_token_loss路径。具体的实现逻辑集中在 optimizer-window metadata provider 中:
prepare_policy_optimizer_window_metadata在每次参数更新前统计全局__dr_grpo_window_scale__。MegatronTrainRayActor将该 scale 绑定到当前 optimizer window;bind_optimizer_window_metadata和DataIterator.get_next负责将其传递给对应的 micro-batch。loss_function在 loss 计算中应用该 scale,并继续复用 Megatron 原有的 gradient accumulation 和 token normalization。这种设计将 Dr.GRPO 的 optimizer-window 统计与通用的 micro-batch 迭代逻辑解耦:算法侧负责生成 scale,data iterator 只负责传递,不需要理解 Dr.GRPO 或 metadata 字段的算法含义。因此,micro-batch 的切分方式不会改变同一 optimizer window 的 loss scale,通用 data iterator 中也不需要增加 Dr.GRPO 分支。
有效 token 计数
get_cp_local_num_tokens去掉cp_size == 1分支对每条 response 的clamp_min(., 1),与
cp_size > 1统一为同一套真实有效 token 定义,使 finalizer 的分母与T = 0的 optimizer window,将跳过 optimizer 和 scheduler,同时报告 warning 信息。T = 0时将报告 0,不再出现除零错误;被跳过的窗口grad_norm保持0.0而非NaN,避免被其他检查误判。Testing
单元测试
Dr.GRPO 测试覆盖:
cp_size = 1/2/4下一致为 0。--normalize-advantages与 Dr.GRPO 组合时被拒绝。train_one_step在空窗口下不调用 optimizer 与 scheduler 的.step(),且grad_norm为0.0而非NaN;正常窗口下两者各调用一次。finalize_model_grads的1/T缩放落在sum(loss) / (N * B),且T == 0时不执行缩放、梯度保持有限。Mixed-length reduction oracle
tests/backends/megatron/test_dr_grpo.py用真实 Gloo process group 把该窗口的数值断言固化为仓库内可执行测试。
窗口为 16 条 response,长度
1024/256/64/16各 4 条,配不规则 prompt 长度与三种loss mask 形态,并按长度 block 连续分片使 DP rank 间的 token 负载不均衡。具体而言:
而六种
(DP, CP)拓扑 × 两种 micro-batch size 下,Dr.GRPO 的实测 loss 与上表完全一致。同一份窗口还在每个拓扑上顺序覆盖:
tp * cp * 2对齐布局与普通 THD 隐式 padding命中相同的
T与 loss,padding token 不进入任何分子或分母。reduction,组合公式成立。
T,clip fraction 与 PG 分子命中冻结值。73x的长度跨度。T保持一致,finalizer 之后的梯度不随 CP degree 变化。
并冻结以下 contract:
(N, T)只在 DP-without-CP group 上归约,不重复统计 CP replica。rollout_max_response_len时 Dr.GRPO 结果按比例变化、GRPO 不变,以区分 configured budget 与 observed max response length。
测试结果
Qwen3.5-4B 200-step end-to-end comparison
使用 Qwen3.5-4B、4 张 H20、GSM8K 和全参数更新,对 GRPO 与 Dr.GRPO 进行 paired comparison。两者使用同一模型、数据、sampling、batch、seed、optimizer 和 nominal learning rate,只切换
ADVANTAGE_ESTIMATOR。具体的复现信息见docs/public/dr-grpo/。Outcome metrics:
Gradient:
Policy-reference KL
train/kl_loss是各自 objective 的归一化分量,不能跨算法比较,因此改用统一定义sum(KL) / T呈现两组的 policy-reference KL。在统一口径下,Dr.GRPO 的 policy-reference KL 高于 GRPO,与其 response 长度显著缩短一致;两组的准确率相当。本实验kl_loss_coef = 0,KL 未参与优化目标,该差异反映的是相对 reference 的分布位移幅度,不构成算法优劣结论。Current-head CP2 与全零 mask smoke
我们进行了
CP=2的 smoke test,覆盖正常、全零和半置零三种 loss mask。三组均正常完成10/10steps,指标全部 finite,无 OOM、collective hang 或 Ray job failure。全零 mask 的窗口按预期跳过 optimizer 与 scheduler,grad_norm为0.0而非NaN;半置零窗口不触发该分支,optimizer 正常执行。pre-commit run --all-filespassespytest tests/)Type of Change