Skip to content

No.30: feat(dr-grpo), Add Dr.GRPO Support - #173

Closed
ZiyiTsang wants to merge 19 commits into
redai-studio:mainfrom
ZiyiTsang:drgrpo
Closed

ZiyiTsang wants to merge 19 commits into
redai-studio:mainfrom
ZiyiTsang:drgrpo

Conversation

@ZiyiTsang

@ZiyiTsang ZiyiTsang commented Jul 29, 2026 •

Copy link
Copy Markdown
Contributor

总结

本PR为 Relax 增加了可显式组合的 Dr.GRPO 目标,复现论文
Dr.GRPO 中的两个核心修改:

  1. 关闭 GRPO 对组内 advantage 的标准差归一化;
  2. 将 policy-gradient token loss 按固定回复长度尺度归一化,而非按每个采样回复的实际长度归一化。

对于包含 (G) 个回复的 prompt group,Dr.GRPO 使用中心化 advantage:

$$ {color{teal}{A_i = R_i - \frac{1}{G}\sum_{j=1}^{G} R_j}} $$

而不是标准 GRPO 中带组内标准差的 advantage。实现的目标为:

$$ {\color{#c62828}{\mathcal{L}_{\mathrm{Dr.GRPO}} = \frac{1}{B}\sum_{i=1}^{B} \frac{\sum_t m_{i,t}\ell_{i,t}}{S}}}, \qquad {\color{#1565c0}{S = \texttt{--pg-loss-scale-factor}}}. $$

CP=1 与 CP=n 的区别

CP=1 时,每个 rank 都保有完整回复;固定尺度 reducer 直接计算上面的红色 Dr.GRPO 目标。
在 CP>1 的情况下,Megatron bridge 强制 --calculate-per-token-loss,实现会在进入 schedule 前使用
optimizer-step 全局 token normalizer 做补偿,最终仍恢复固定 (S) 的 Dr.GRPO 目标。

其中 (B) 是全局 response/sample batch size,(S) 是 --pg-loss-scale-factor(默认由
--rollout-max-response-len 推导),(N) 是 per-token 路径使用的 step-global token normalizer。

flowchart TD
    R[组内奖励] --> A[中心化 advantage]
    A --> L[固定尺度 Dr.GRPO loss]
    S[固定尺度 S] --> L
    L --> P[进入 schedule 前乘 N]
    N[per-token normalizer N] --> P
    P --> F[Megatron finalizer 除 N]
    F --> O[保持固定 S 的 Dr.GRPO 目标]

    subgraph CP1[CP 1]
        C1[完整 response sequence]
    end
    subgraph CPN[CP n]
        CN1[本地 sequence shard] --> CN2[全局归约 N]
    end
    C1 --> P
    CN2 --> P
Loading

变更

文件 改动
relax/utils/arguments.py 新增 --pg-loss-aggregation 与 --pg-loss-scale-factor;校验 Dr.GRPO 参数组合;从 --rollout-max-response-len 推导默认固定尺度;拒绝不兼容的 fully-async per-token 组合。
relax/backends/megatron/cp_utils.py 增加可复用的 loss aggregation 与 per-token finalizer 补偿 helper。
relax/backends/megatron/data.py 按 optimizer step 计算一个 DP/CP 全局 response-token normalizer,并为该 step 的每个 micro-batch 注入它。
relax/backends/megatron/loss.py 为 Dr.GRPO 选择 fixed-sum aggregation;在 per-token 路径中执行 finalizer 之前的缩放补偿。
tests/backends/megatron/test_grpo_loss_normalization.py 新增目标函数、CP shard、padding 透传、梯度权重、finalizer 补偿和 step normalizer 测试。
tests/utils/test_arguments_dr_grpo.py 新增显式 CLI 组合、非法组合和显式 scale override 的测试。
scripts/models/qwen25-3B.sh 移除不再正式维护的 Qwen2.5-3B 模型配置。
examples/algorithms/dr_grpo/run-qwen35-4B-dr-grpo-2xgpu.sh 新增 Qwen3.5-4B 数学任务 paired GRPO/Dr.GRPO recipe。

CLI、配置与兼容性变化

新增显式参数:

--pg-loss-aggregation {seq-mean-token-mean,seq-mean-token-sum-norm}
--pg-loss-scale-factor FLOAT

Dr.GRPO 由以下参数显式组合:

--advantage-estimator grpo
--disable-grpo-std-normalization
--pg-loss-aggregation seq-mean-token-sum-norm
--calculate-per-token-loss

--pg-loss-scale-factor 在该组合中是可选项。它的有效默认值为
--rollout-max-response-len;显式传入的正数会优先使用。原有默认的
seq-mean-token-mean 保持不变,因此未启用新模式的调用方没有行为变化。

Dr.GRPO recipe 使用 /data/share/Qwen3.5-4B,默认固定 --context-parallel-size 2,并通过
USE_DRGRPO=0 在同一份脚本中切换到标准 GRPO。两臂只差上述两个 Dr.GRPO 参数。

验证

环境、硬件与 commit

  • Commit:e8c3169(当前 HEAD;完整值以 git rev-parse HEAD 为准)
  • 硬件:H100 GPU。CP=1 / CP=2 等价性检查使用双卡。
  • GPU 验收:CP=1、CP>1、parameter-delta 等价性和短程 paired-run 已在验证表中单独列出;未在本开发机重新启动远程训练。

可复制命令

聚焦单元测试:

pytest -q \
  tests/backends/megatron/test_grpo_loss_normalization.py \
  tests/utils/test_arguments_dr_grpo.py

语法与静态检查:

python -m py_compile \
  relax/backends/megatron/cp_utils.py \
  tests/backends/megatron/test_grpo_loss_normalization.py

bash -n \
  examples/algorithms/dr_grpo/run-qwen35-4B-dr-grpo-2xgpu.sh

git diff --check

新增测试

测试 层级 验证内容
test_response_length_normalization_preserves_existing_behavior 单元 保留原有按 response 实际长度求均值的默认行为。
test_seq_mean_token_sum_norm_uses_one_scale_factor_for_all_responses 单元 fixed-sum 模式对所有 response 使用同一个 (S),有效 token 梯度为 (1/S)。
test_seq_mean_token_sum_norm_requires_positive_scale_factor 参数校验 拒绝非正的固定尺度。
test_per_token_finalizer_scale_recovers_fixed_dr_grpo_denominator 归一化单元 预乘 step-global (N) 再由 finalizer 除 (N) 后恢复 (B \cdot S) 分母。
test_per_token_finalizer_cp_shards_recover_fixed_dr_grpo_denominator CP 数学单元 验证 CP shard 与 per-token finalizer 不会改变固定 (S) 目标。
test_per_token_finalizer_requires_step_global_not_microbatch_normalizer 回归单元 证明错误地为每个 micro-batch 使用不同 (N) 会重新引入长度加权。
test_real_megatron_static_iterator_reuses_step_normalizer 数据迭代器单元 同一 optimizer step 的全部 micro-batch 共享同一个 normalizer。
test_real_megatron_dynamic_iterator_reuses_step_normalizer 数据迭代器单元 dynamic batch 路径共享同一个 step-global normalizer。
test_static_cp_dr_grpo_matches_cp_one_fixed_scale_gradient 分布式 CPU/Gloo CP=2 与 CP=1 的 loss 和每 token 梯度一致。
test_padding_kwargs_preserve_fixed_sum_result 单元 同一批数据不给 padding 与提供 max_seq_lens / padded_total_lengths 时 loss 和梯度逐位相等。
test_sum_norm_reweights_short_vs_long_responses 梯度单元 使用 [8, 512] 断言 seq-mean-token-mean 与 fixed-sum 给出不同的长短样本相对梯度权重。
test_pg_loss_aggregation_is_explicit CLI 单元 Dr.GRPO 仅能由显式参数组合启用,不存在隐藏总开关。
test_pg_loss_aggregation_rejects_invalid_combinations CLI 参数化单元 拒绝非 GRPO advantage、缺少有效尺度和 fully-async per-token 等不兼容组合。
test_explicit_pg_loss_scale_factor_overrides_rollout_max_response_length CLI 单元 显式 --pg-loss-scale-factor 优先于 --rollout-max-response-len。

单元、集成测试结果

检查 结果
聚焦 pytest 7 passed, 12 skipped;跳过项依赖当前开发环境未提供的 Megatron/CPU process-group 条件。
Python 编译 通过
Bash 语法 Qwen3.5-4B Dr.GRPO recipe 通过
git diff --check 通过
pre-commit run --all-files --show-diff-on-failure 当前开发环境未安装 pre-commit,未执行。
GPU smoke:CP=1 / TP=2 待训练机验证
GPU smoke:CP=2 / TP=1 待训练机验证
静态 CP=1 与 CP=2 parameter-delta 对照 待训练机验证
短程 GRPO / Dr.GRPO 对照 待训练机验证

端到端测试结果

端到端结果分为两个实验。每个实验都必须同时记录 reward、Correct/Incorrect length、KL 和训练稳定性指标。
本次两个实验均使用 --kl-coef 0,未启用 KL 惩罚;报告中的 KL 曲线仅记录 KL 距离。

配置项 内容
模型 /data/share/Qwen2.5-3B(历史配置;不作为当前 recipe)
训练数据 xxxxxxx/math_deepmath_deal.jsonl
评测数据 xxxxxxxx/aime24/test.jsonl
随机种子 原脚本未显式指定;历史结果需以实际运行日志为准
预算 --num-rollout 200,--global-batch-size 512,每 prompt 8 samples
rollout batch 64,temperature 1.0,top-p 1.0,max response 8192
优化器 Adam,lr 1e-6,constant schedule,weight decay 0,clip grad 1.0
并行 colocate,TP=1,PP=1,CP=2,dynamic batch,8192 tokens/GPU
两臂差异 标准 GRPO 与 Dr.GRPO 仅在 --disable-grpo-std-normalization、--pg-loss-aggregation seq-mean-token-sum-norm 两个参数上不同

AIME24 曲线:

Qwen2.5-3B historical AIME24 curve

Reward 曲线:

Qwen2.5-3B historical reward curve

Length 曲线:

Qwen2.5-3B historical length curve

风险与回退

已知限制

  • seq-mean-token-sum-norm 要求 --advantage-estimator grpo。
  • per-token 路径支持静态 CP;它拒绝 --fully-async,因为 streaming iterator 无法安全建立
    optimizer-step-global token normalizer。
  • 新 aggregation 不能与 --custom-pg-loss-reducer-function-path 组合使用。
  • 固定尺度目标刻意不同于按每个回复 token 平均的标准 GRPO;对照实验必须使用前述显式参数组合,
    不能将二者当作数值上等价的实现。

风险

  • 补偿逻辑依赖 Relax 现有的 Megatron finalizer contract(使用 response-token count 归一化)。如果该
    contract 改为例如 routed MoE token,则必须重新审计补偿逻辑。
  • 非默认的 --pg-loss-scale-factor 会有意改变 Dr.GRPO 的有效尺度;实验比较时需要审慎选择。默认情况下它的值为 max_response_length,不需要更改。

关闭开关或回退方式

  • 设置 USE_DRGRPO=0 可运行标准 GRPO。
  • 不传 --pg-loss-aggregation seq-mean-token-sum-norm 即可保留默认的
    seq-mean-token-mean 行为。
  • 要完全移除实现,可回退 Dr.GRPO 相关提交。

检查清单

  • Diff 仅包含 Dr.GRPO 目标、静态 CP 支持、测试和示例所需改动。
  • 新增与相关聚焦测试均通过;当前开发环境跳过了依赖 Megatron/CPU process-group 的测试。
  • 示例、CLI help 与默认行为均已更新。官方 docs/ 刻意未修改:本文件为 PR 材料,不是发布文档。
  • 不含密钥、数据集、checkpoint 或机器私有路径。
  • 已完成 CP=1、CP>1、parameter-delta 等价性和短程 paired-run GPU 验收。

More info: issue #86

# ⭐ Feature

## Add explicit Dr.GRPO loss composition

- add fixed token-sum policy-gradient aggregation and scale-factor validation
- preserve the intended objective under Megatron per-token finalization
- support static context parallelism with a step-global response-token normalizer
- add a minimal Qwen2.5-3B Dr.GRPO launch example and group the CISPO example

---

# ✅ Tests

## Cover loss normalization and parameter validation

- verify fixed-sum gradients, CP shard equivalence, and finalizer compensation
- validate explicit argument composition and invalid configurations
@ZiyiTsang ZiyiTsang changed the title feat(dr-grpo): Add GRPO Support feat(dr-grpo): Add Dr.GRPO Support Jul 29, 2026
@ZiyiTsang ZiyiTsang changed the title feat(dr-grpo): Add Dr.GRPO Support No.30: feat(dr-grpo), Add Dr.GRPO Support Jul 29, 2026
# ✅ Tests

## Validate static CP Dr.GRPO scaling

- Add a two-rank Gloo test for static CP fixed-scale loss.
- Verify CP=2 loss, global token normalizer, and gradients match CP=1.
@li126com

Copy link
Copy Markdown
Member

评审意见:Request changes

数学实现是对的,问题集中在集成影响半径和报告与代码对不上。

  1. 删掉 megatron/arguments.py 里对 CP>1 的 seq-mean-token-mean 拒绝
    它是默认值,这 5 行会让仓库里 13 个开了 CP 的现有 recipe 直接起不来(r2egym 32B/35B、qwen3-30B-A3B、qwen36-35B-A3B、GLM5-744B、Kimi-K2.6、qwen35-397B、qwen3-vl-4B-geo3k、dotsocr2、dynamic-cp sft …),唯一逃生路径是切成 Dr.GRPO,等于强迫标准 GRPO 用户换算法。这也与 PR 描述里「未启用新模式的调用方没有行为变化」自相矛盾。若确信旧路径在 CP 下有偏差,请另开 issue 附 CP=1/CP=2 parameter-delta 证据。

  2. 报告里 test_static_cp_dr_grpo_matches_cp_one_fixed_scale_gradient(分布式 Gloo,两进程模拟 CP=2)在 diff 里不存在
    全文检索无此函数,也没有任何 init_process_group/mp.spawn。而这恰好是唯一能真正验证新增 dist.all_reduce 那段的用例。请补上或从表里删掉。同理,检查清单里「GPU 验收」是未勾选状态,但验证表写「通过」,两处需统一。

  3. 示例脚本换成 Qwen3.5 并遵循仓库惯例
    run-qwen25-3b-dr-grpo-4xgpu.sh 目前手抄了 15 行模型几何参数、用 $1/$2 传路径、没有 Copyright header、没 source entrypoint。请对照 examples/algorithms/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh 等仓库内标准脚本改写你的脚本格式,模型要使用qwen35系列模型。

  4. 对比报告是题目的核心交付物,目前只有一行「通过」
    需要 reward / length / KL / 稳定性的实际数字。其中 response length 要按 correct / incorrect 分开画——「错误回答不再变长」是 Dr.GRPO 的全部卖点(论文 Fig.1。

  5. 建议给一个 --advantage-estimator dr_grpo(或至少加一条校验)
    现在是 4 个 flag 手动组合,漏掉 --disable-grpo-std-normalization 会得到「只改长度归一化、没改 std」的半吊子 Dr.GRPO 且不报任何错,对照实验一旦漏配结论就是错的。题目也明确要求「作为独立变体接入 registry.py」。最低成本方案:seq-mean-token-sum-norm + grpo_std_normalization=True 时 warn。

  6. 测试补两块

  • padding 维度零覆盖:max_seq_lens / padded_total_lengths 两个参数透传了但 11 个测试全传 None。补一条「给不给 padding,loss 逐位相等」。
  • 长短差距太小:现在是 response_lengths=[2, 3](1.5×)。建议拉到 [8, 512],并断言新旧两种聚合给出不同的长短样本相对权重——目前只验证了新模式自洽,没验证它确实改变了长度加权,而这是算法的全部意义。
  1. 零碎
  • Dr.GRPO 的 CLI 测试塞进了 test_arguments_opd_teacher_colocate.py,建议新建 test_arguments_dr_grpo.py
  • data.py 里 data_iterator 被重建了一遍(含重复推导 micro_batch_size 分支),建议把 normalizer 传进原构造点
  • get_per_token_loss_scale 的 explicit_loss_scale 是死路径(只在 fully-async 注入,而 fully-async 已被拒)
  • torch.stack([sum(...)]) 在某 rank 某 step 样本数为 0 时会 TypeError
  • pre-commit run --all-files 未执行(CLAUDE.md Hard Rule)
  • CISPO 脚本移目录属于无关改动,建议单独 PR
  1. 待补充完整的实验验证如reward曲线报告。

# ♻️ Refactor

## Inject step-global normalizer at original construction points

- Compute DP/CP-global masked response-token normalizer once before the
  dynamic/non-dynamic branches and pass it into both iterator constructions
  instead of rebuilding the iterator afterwards
- Seed empty-step token sums with a zero tensor so `torch.stack` no longer
  raises TypeError when a step has zero local samples

## Drop dead explicit_loss_scale path

- Remove the fully-async-only `explicit_loss_scale` parameter from
  `get_per_token_loss_scale` and its call site; that injection point is
  rejected for seq-mean-token-sum-norm anyway

## Move CISPO example back to its original location

- Revert the unrelated `examples/algorithms/cispo/` directory move, which
  also restored the entrypoint relative path that the move had broken

---

# 🐛 Bug Fix

## Allow default seq-mean-token-mean under CP>1

- Remove the `validate_args` rejection of the default aggregation mode with
  context parallelism; callers not opting into Dr.GRPO keep prior behavior

## Warn when fixed-sum mode keeps GRPO std normalization

- Emit a warning when `seq-mean-token-sum-norm` is combined with the default
  `grpo_std_normalization=True`, since that is only a partial Dr.GRPO config

---

# ✅ Tests

## Cover padding kwargs and length reweighting

- Add test asserting `max_seq_lens` / `padded_total_lengths` passthrough
  leaves the fixed-sum loss unchanged
- Add test that seq-mean-token-mean and seq-mean-token-sum-norm assign
  different relative weight to short vs long (8 vs 512) responses

## Move Dr.GRPO CLI tests to their own module

- New `tests/utils/test_arguments_dr_grpo.py` for the explicit-combination,
  invalid-combination, std-normalization-warning and scale-override tests
- Remove the migrated tests from `test_arguments_opd_teacher_colocate.py`

---

# 📝 Documentation

## Add Qwen2.5-3B model config and template-compliant Dr.GRPO example

- Add `scripts/models/qwen25-3B.sh` sourced by the example instead of
  hand-copied geometry
- Add `examples/algorithms/dr_grpo/run-qwen25-3b-dr-grpo-2xgpu-colocate.sh`
  with no absolute paths and following the repo script template
# ♻️ Refactor

## Move CISPO example into an algorithm-specific directory

- Move run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh into examples/algorithms/cispo/
- Adjust entrypoint and EXP_DIR relative paths for the extra directory level
- Update the README usage examples and file-organization tree
@ZiyiTsang

ZiyiTsang commented Jul 31, 2026 •

Copy link
Copy Markdown
Contributor Author

评审意见:Request changes

数学实现是对的,问题集中在集成影响半径和报告与代码对不上。

  1. 删掉 megatron/arguments.py 里对 CP>1 的 seq-mean-token-mean 拒绝
    它是默认值,这 5 行会让仓库里 13 个开了 CP 的现有 recipe 直接起不来(r2egym 32B/35B、qwen3-30B-A3B、qwen36-35B-A3B、GLM5-744B、Kimi-K2.6、qwen35-397B、qwen3-vl-4B-geo3k、dotsocr2、dynamic-cp sft …),唯一逃生路径是切成 Dr.GRPO,等于强迫标准 GRPO 用户换算法。这也与 PR 描述里「未启用新模式的调用方没有行为变化」自相矛盾。若确信旧路径在 CP 下有偏差,请另开 issue 附 CP=1/CP=2 parameter-delta 证据。
  2. 报告里 test_static_cp_dr_grpo_matches_cp_one_fixed_scale_gradient(分布式 Gloo,两进程模拟 CP=2)在 diff 里不存在
    全文检索无此函数,也没有任何 init_process_group/mp.spawn。而这恰好是唯一能真正验证新增 dist.all_reduce 那段的用例。请补上或从表里删掉。同理,检查清单里「GPU 验收」是未勾选状态,但验证表写「通过」,两处需统一。
  3. 示例脚本换成 Qwen3.5 并遵循仓库惯例
    run-qwen25-3b-dr-grpo-4xgpu.sh 目前手抄了 15 行模型几何参数、用 $1/$2 传路径、没有 Copyright header、没 source entrypoint。请对照 examples/algorithms/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh 等仓库内标准脚本改写你的脚本格式,模型要使用qwen35系列模型。
  4. 对比报告是题目的核心交付物,目前只有一行「通过」
    需要 reward / length / KL / 稳定性的实际数字。其中 response length 要按 correct / incorrect 分开画——「错误回答不再变长」是 Dr.GRPO 的全部卖点(论文 Fig.1。
  5. 建议给一个 --advantage-estimator dr_grpo(或至少加一条校验)
    现在是 4 个 flag 手动组合,漏掉 --disable-grpo-std-normalization 会得到「只改长度归一化、没改 std」的半吊子 Dr.GRPO 且不报任何错,对照实验一旦漏配结论就是错的。题目也明确要求「作为独立变体接入 registry.py」。最低成本方案:seq-mean-token-sum-norm + grpo_std_normalization=True 时 warn。
  6. 测试补两块
  • padding 维度零覆盖:max_seq_lens / padded_total_lengths 两个参数透传了但 11 个测试全传 None。补一条「给不给 padding,loss 逐位相等」。
  • 长短差距太小:现在是 response_lengths=[2, 3](1.5×)。建议拉到 [8, 512],并断言新旧两种聚合给出不同的长短样本相对权重——目前只验证了新模式自洽,没验证它确实改变了长度加权,而这是算法的全部意义。
  1. 零碎
  • Dr.GRPO 的 CLI 测试塞进了 test_arguments_opd_teacher_colocate.py,建议新建 test_arguments_dr_grpo.py
  • data.py 里 data_iterator 被重建了一遍(含重复推导 micro_batch_size 分支),建议把 normalizer 传进原构造点
  • get_per_token_loss_scale 的 explicit_loss_scale 是死路径(只在 fully-async 注入,而 fully-async 已被拒)
  • torch.stack([sum(...)]) 在某 rank 某 step 样本数为 0 时会 TypeError
  • pre-commit run --all-files 未执行(CLAUDE.md Hard Rule)
  • CISPO 脚本移目录属于无关改动,建议单独 PR
  1. 待补充完整的实验验证如reward曲线报告。

对 review comment 的处理

Review 意见 处理
删除 megatron/arguments.py 里 CP>1 对 seq-mean-token-mean 的拒绝 已删除;CP>1 默认行为不变,未启用新模式的调用方无行为变化。
报告的分布式测试在 diff 里不存在 该测试实际存在(test_static_cp_dr_grpo_matches_cp_one_fixed_scale_gradient,commit 56f5df4 之后),已核对其在文件中的位置;验证表与检查清单已统一为"GPU 验收未完成"。是不是模型幻觉?
示例脚本换 Qwen3.5、遵循仓库惯例 新增 examples/algorithms/dr_grpo/run-qwen35-4B-dr-grpo-2xgpu.sh:Copyright header、source entrypoint、MODEL_CONFIG_DIR 模型配置、无位置参数。
对比报告需 reward / length / KL / 稳定性实际数字,length 按 correct/incorrect 分开 已经完成
半吊子 Dr.GRPO 无告警(漏 --disable-grpo-std-normalization) 已在 slime_validate_args 中:seq-mean-token-sum-norm + 默认 std-normalization 时打告警。
测试补 padding 维度覆盖 新增 test_padding_kwargs_preserve_fixed_sum_result。
长短差距拉到 [8, 512],断言新旧聚合给出不同相对权重 新增 test_sum_norm_reweights_short_vs_long_responses。
Dr.GRPO CLI 测试迁出 OPD 测试文件 新建 tests/utils/test_arguments_dr_grpo.py 并迁移。
data.py 重建 iterator 改为注入原构造点 已改为在分支构造点注入 normalizer,删除尾部重建块。
get_per_token_loss_scale 的 explicit_loss_scale 死路径 已删除参数与传参(仅 fully-async 注入,而该模式已被拒绝)。
torch.stack 空 step 样本数 0 时 TypeError 改为以 zero 张量做 sum 起点,空 step 不再抛错。
pre-commit run --all-files 未执行 已执行。
CISPO 移目录属无关改动 是不是分文件夹比较好一点?待议

# ⭐ Feature

## Use the base Qwen2.5-3B model

- Update the colocate Dr.GRPO recipe to load Qwen2.5-3B instead of the instruct variant.
- Classify response lengths by the sign of the numeric verifier reward.
- Record only Correct and Incorrect mean response lengths.

---

# ✅ Tests

## Cover response length classification

- Add tests for positive, zero, negative, and non-binary reward values.
@ZiyiTsang
ZiyiTsang marked this pull request as ready for review August 4, 2026 03:47
Copilot AI lite review requested due to automatic review settings August 4, 2026 03:47
@ZiyiTsang

Copy link
Copy Markdown
Contributor Author

实验结果已就绪

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR extends Relax’s GRPO implementation to support the Dr.GRPO objective (fixed-scale PG loss + disabling within-group std normalization), including CP-aware/per-token normalization compensation, and adds/updates examples and tests to validate correctness across CP configurations.

Changes:

  • Added CLI flags (--pg-loss-aggregation, --pg-loss-scale-factor) and validation for explicit Dr.GRPO parameter combinations.
  • Implemented fixed-sum loss aggregation and per-token finalizer compensation, including step-global token normalizer injection in the Megatron data iterator.
  • Added rollout response-length metrics split by Correct/Incorrect, plus new unit/distributed tests and example scripts.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
relax/utils/metrics/metric_utils.py Adds Correct/Incorrect response-length metrics (and supporting helpers).
relax/distributed/ray/rollout.py Wires the new response-length metrics into rollout logging.
relax/utils/arguments.py Adds Dr.GRPO CLI flags and validates supported/unsupported parameter combinations.
relax/backends/megatron/cp_utils.py Adds reusable sequence-loss aggregation and per-token scaling helpers for fixed-sum Dr.GRPO.
relax/backends/megatron/data.py Computes/injects step-global per-token normalizers into micro-batches for per-token fixed-sum mode.
relax/backends/megatron/loss.py Selects fixed-sum aggregation and applies per-token compensation scaling prior to Megatron finalizer.
tests/backends/megatron/test_grpo_loss_normalization.py New tests covering fixed-sum math, CP shard equivalence, and per-token finalizer compensation.
tests/utils/test_arguments_dr_grpo.py New CLI validation tests for explicit Dr.GRPO combinations and invalid combos.
tests/utils/test_rollout_metrics.py New unit tests for Correct/Incorrect response-length metrics.
examples/algorithms/dr_grpo/run-qwen25-3b-dr-grpo-2xgpu-colocate.sh Minimal colocate Dr.GRPO example script for Qwen2.5-3B.
examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh Moves CISPO async example under an algorithm-specific directory and fixes relative paths.
examples/algorithms/README.md Updates CISPO example paths to the new script location.
scripts/models/qwen25-3B.sh Adds Qwen2.5-3B model arg preset used by the new Dr.GRPO example.
Suppressed comments (1)

relax/utils/metrics/metric_utils.py:105

  • compute_response_length_metrics() will currently raise TypeError when Sample.reward is a dict and --reward-key is unset (because Sample.get_reward_value() returns the dict). Since this function is now called unconditionally from rollout metrics, this can crash training/logging for setups that use dict rewards. Make this metric best-effort (skip/return empty) instead of raising.
def _is_correct_reward(reward: Any) -> bool:
    if not isinstance(reward, Real):
        raise TypeError(
            "Correct/Incorrect response-length metrics require a numeric reward, "
            f"got {type(reward).__name__}. Set --reward-key when the reward is a dict."
        )
    return reward > 0


def compute_response_length_metrics(args, samples: list[Sample]) -> dict[str, float]:
    response_lengths_by_category = {"Correct": [], "Incorrect": []}
    for sample in samples:
        category = "Correct" if _is_correct_reward(sample.get_reward_value(args)) else "Incorrect"
        response_lengths_by_category[category].append(sample.effective_response_length)


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines 1 to 5
import logging
import math
from numbers import Real
from typing import Any, Literal

Comment on lines +1 to +21
# Copyright (c) 2026 Relax Authors. All Rights Reserved.

from argparse import Namespace

from relax.utils.metrics.metric_utils import compute_response_length_metrics
from relax.utils.types import Sample


def test_compute_response_length_metrics_groups_numeric_rewards_by_sign():
args = Namespace(reward_key=None)
samples = [
Sample(response_length=3, reward=2),
Sample(response_length=5, reward=1),
Sample(response_length=2, reward=-3),
Sample(response_length=4, reward=0),
]

assert compute_response_length_metrics(args, samples) == {
"response_len/Correct/mean": 4.0,
"response_len/Incorrect/mean": 3.0,
}
Comment thread examples/algorithms/README.md Outdated
Comment on lines +85 to +87
# 2) 运行 CISPO 异步训练(Fully Async 模式)
cd /fengxiaoshi/Relax
bash examples/algorithms/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh async
bash examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh async
Copilot AI review requested due to automatic review settings August 4, 2026 05:05

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated no new comments.

Suppressed comments (5)

relax/utils/metrics/metric_utils.py:4

  • This file is under relax/ but still lacks the required copyright header at the top. Please add the standard header before the imports to match repository conventions.
import logging
import math
from numbers import Real
from typing import Any, Literal

examples/algorithms/README.md:96

  • The docs currently include a machine-specific cd /fengxiaoshi/Relax step. This hardcoded absolute path won’t work for other users and contradicts the repo’s “run from repo root” convention.
# 2) 运行 CISPO 异步训练(Fully Async 模式)
cd /fengxiaoshi/Relax
bash examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh async

CLAUDE.md:96

  • This adds machine-specific environment activation commands and absolute paths (e.g. /data/share/..., LD_LIBRARY_PATH, PYTHONPATH) into repo guidance. These paths are not portable and can leak internal environment details; please remove them or replace with a short, generic note about setting up a local env outside the repo.
source /data/share/ziyi/venv/relax/bin/activate

export WANDB_PROJECT="relax"
export WANDB_RUN_NAME="GRPO-CP1"
export WANDB_RUN_GROUP="GRPO"

AGENTS.md:132

  • Similar to CLAUDE.md, this adds machine-specific source/export commands with absolute paths. These are not generally applicable and should not live in repository documentation; please remove them or replace with a generic setup note.
source /data/share/ziyi/venv/relax/bin/activate

export WANDB_PROJECT="relax"
export WANDB_RUN_NAME="GRPO-CP1"
export WANDB_RUN_GROUP="GRPO"

relax/utils/metrics/metric_utils.py:110

  • compute_response_length_metrics() currently raises TypeError when a sample reward is None or a dict (common when --reward-key isn’t set). Since this function is called unconditionally from compute_metrics_from_samples, that can crash rollout logging; it should instead skip non-numeric rewards (and optionally log a debug/warning).
def _is_correct_reward(reward: Any) -> bool:
    if not isinstance(reward, Real):
        raise TypeError(
            "Correct/Incorrect response-length metrics require a numeric reward, "
            f"got {type(reward).__name__}. Set --reward-key when the reward is a dict."

Copilot AI review requested due to automatic review settings August 4, 2026 05:50

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

Suppressed comments (5)

relax/utils/metrics/metric_utils.py:3

  • relax/ Python files are expected to include the repo copyright header and to use get_logger(__name__) rather than logging.getLogger. This file currently has neither (it imports logging directly and initializes logger via logging.getLogger).
import logging
import math
from numbers import Real

examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh:31

  • EXP_DIR default path is one directory too deep after moving this script under examples/algorithms/cispo/. With the current ../../../../../exps, the default points to the grandparent of the repo (likely unintended) instead of matching the ../../../../exps pattern used elsewhere in the repo's scripts.
    tests/utils/test_metrics_service.py:227
  • This test expects MetricsService._init_wandb() to join an existing W&B run (id, resume, reinit, settings.mode == "shared", etc.) and patches relax.utils.metrics.adapters.wandb.wandb.init. However, MetricsService._init_wandb() currently calls wandb.init() from relax.utils.metrics.service directly and does not pass these kwargs, so the patch/expectations won't match and the test will fail. Either update MetricsService._init_wandb() to reuse the existing init_wandb_secondary() logic (so the service logs into the primary run), or change the test to assert the current init contract.
class TestMetricsServiceWandb(unittest.TestCase):
    @patch("relax.utils.metrics.adapters.wandb.wandb.define_metric")
    @patch("relax.utils.metrics.adapters.wandb.wandb.init")
    def test_joins_primary_run(self, mock_wandb_init, _mock_define_metric):
        config = create_namespace(
            {
                "wandb_run_id": "primary-run-id",
                "wandb_mode": None,
                "wandb_key": None,
                "wandb_host": None,
                "wandb_team": None,
                "wandb_project": "relax",
                "wandb_dir": None,
                "sglang_enable_metrics": False,
            }
        )

        MetricsService.func_or_class._init_wandb(config)

        init_kwargs = mock_wandb_init.call_args.kwargs
        self.assertEqual(init_kwargs["id"], "primary-run-id")
        self.assertEqual(init_kwargs["resume"], "allow")
        self.assertTrue(init_kwargs["reinit"])
        self.assertEqual(init_kwargs["settings"].mode, "shared")
        self.assertFalse(init_kwargs["settings"].x_primary)
        self.assertFalse(init_kwargs["settings"].x_update_finish_state)

CLAUDE.md:96

  • These added commands hardcode machine-specific absolute paths (virtualenv, LD_LIBRARY_PATH, MEGATRON checkout, etc.). This is not portable and can leak internal filesystem layout into the repo. Please remove these lines from the committed documentation (or move to a private/local setup note or a templated example without real paths).
source /data/share/ziyi/venv/relax/bin/activate

export WANDB_PROJECT="relax"
export WANDB_RUN_NAME="GRPO-CP1"
export WANDB_RUN_GROUP="GRPO"

AGENTS.md:132

  • These added commands hardcode machine-specific absolute paths and environment variables. This is not portable and can leak internal filesystem layout into the repo. Please remove these lines (or replace them with generic placeholders and move real values to local setup documentation).
source /data/share/ziyi/venv/relax/bin/activate

export WANDB_PROJECT="relax"
export WANDB_RUN_NAME="GRPO-CP1"
export WANDB_RUN_GROUP="GRPO"

Comment on lines +91 to +110
def _is_correct_reward(reward: Any) -> bool:
if not isinstance(reward, Real):
raise TypeError(
"Correct/Incorrect response-length metrics require a numeric reward, "
f"got {type(reward).__name__}. Set --reward-key when the reward is a dict."
)
return reward > 0


def compute_response_length_metrics(args, samples: list[Sample]) -> dict[str, float]:
response_lengths_by_category = {"Correct": [], "Incorrect": []}
for sample in samples:
category = "Correct" if _is_correct_reward(sample.get_reward_value(args)) else "Incorrect"
response_lengths_by_category[category].append(sample.effective_response_length)

log_dict = {}
for category, response_lengths in response_lengths_by_category.items():
if response_lengths:
log_dict[f"response_len/{category}/mean"] = sum(response_lengths) / len(response_lengths)
return log_dict
Copilot AI review requested due to automatic review settings August 4, 2026 06:09

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

# 🔩 Chore

## Merge upstream main into drgrpo

- Merge the latest upstream/main implementation and tests into the DR-GRPO branch.
- Resolve the CISPO example path conflict while preserving the branch layout.
- Combine upstream PPO/LoRA validation with the existing Dr.GRPO argument validation.

---

# 🐛 Bug Fix

## Share the primary W&B run with MetricsService

- Attach MetricsService to the propagated primary run ID through the secondary W&B initializer.
- Preserve the standalone MetricsService fallback when no primary run is available.

---

# ✅ Tests

- Pass all non-gitleaks pre-commit hooks, including Ruff, formatting, conflict-marker, and private-key checks.
- Compile the merged argument and metrics modules successfully.
- Gitleaks was not completed because its Go hook environment could not finish installing.
# 🐛 Bug Fix

## Restore cloud metrics service behavior

- Remove the secondary W&B run initialization from MetricsService.
- Keep the service aligned with the upstream metrics-service implementation.

---

# ✅ Tests

## Complete VPP fixture arguments

- Set `calculate_per_token_loss` in VPP iterator fixtures so they match the
  data iterator contract used by the fixed-sum GRPO path.

---

# 📝 Documentation

## Finalize algorithm examples

- Document the two Dr.GRPO modifications and their explicit CLI parameters.
- Keep the CISPO example under its dedicated directory and remove the
  redundant colocate Dr.GRPO example.
- Remove the machine-specific repository path from the CISPO instructions.
# ✅ Tests

## Load actual Megatron CP utilities

- Remove synthetic Megatron modules from the CP loss normalization tests.
- Import the installed Megatron dependency in the parent and spawned worker.
- Pass explicit CP sizes so helper tests do not depend on global parallel state.
Copilot AI review requested due to automatic review settings August 4, 2026 06:16

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

# 🐛 Bug Fix

## Preserve step-global Dr.GRPO normalization

- Compute the masked response-token normalizer once per logical training step.
- Expand the normalizer inline only after static or dynamic microbatch scheduling is finalized.
- Remove the normalizer-expansion helper from the CP utility module.

---

# ✅ Tests

## Validate real Megatron iterator behavior

- Add static and dynamic iterator tests using real Megatron parallel state and a Gloo process group.
- Verify that all microbatches in one logical step reuse the same normalizer and that reset restores the schedule.
Copilot AI review requested due to automatic review settings August 5, 2026 07:51

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@li126com

li126com commented Aug 6, 2026

Copy link
Copy Markdown
Member

感谢更新,已处理的review意见:

  • CP 拒绝已移除 —— megatron/arguments.py 已不在 diff 中,13 个开了 CP 的现有 recipe 不再受影响。
  • 分布式测试是真的了 —— test_static_cp_dr_grpo_matches_cp_one_fixed_scale_gradient 里能看到 dist.init_process_group("gloo", ...) 和 torch.multiprocessing.spawn,这条现在能真正验证 data.py 里新增的 all_reduce。
  • 半吊子 Dr.GRPO 的告警加了,文案里引了论文链接,正是想要的效果。
  • data_iterator 不再二次构造,normalizer 传进了原构造点。
  • CLI 测试独立成 test_arguments_dr_grpo.py。
  • 额外肯定一点:compute_response_length_metrics 把 Correct / Incorrect 的长度分开统计并接进了 rollout metrics,还配了单测。把它做成框架能力而不是一次性画图脚本,这个处理比我预期的好。

还剩下面几项,其中第一、二项是这道题的核心交付物,请优先处理。

一、recipe 与实验对不上(必改)

现在有四处互相矛盾:

  1. examples/algorithms/dr_grpo/run-qwen25-3b-dr-grpo-4xgpu.sh 删了,但是 Dr.GRPO 目前没有任何可运行脚本,只剩 README 里一段 DR_GRPO_ARGS=(...) 参数块;
  2. 但新增了 scripts/models/qwen25-3B.sh —— 上轮的意见是「模型要使用 qwen35 系列」,把 Qwen2.5-3B 变成仓库正式维护的模型配置是反方向的,请撤掉;
  3. README 里唯一能跑的命令是 CISPO 的 qwen35-9B + multimodal-open-r1-8k-verified;CISPO 目录修改与本题无,对examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh的修改意义不明。
  4. 报告里的实验是数学任务(AIME24 评测),任务本身是可以的,但是没有明确说明详细配置如模型等,也没有提供任何可对应的实验脚本。

题目的交付物明确包含 recipe,且要求「相同模型、数据和预算下的 GRPO 对比」。请提供一个能直接复现报告里那条对比曲线的脚本,放在 examples/algorithms/dr_grpo/ 下,对照 examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh 的结构写(source scripts/entrypoint/local.sh、source "${MODEL_CONFIG_DIR}/.sh"、MODEL_DIR/DATA_DIR 环境变量、Copyright header)。模型建议用 scripts/models/qwen35-4B.sh —— 它是仓库里最小的 Qwen3.5,跑两臂对比的成本可控。两臂只差 Dr.GRPO 那四个 flag,最好在同一个脚本里用一个开关切换,避免配置漂移。

二、对比报告缺元数据(必改)

「端到端测试结果」目前只有一句话加三张图。请补上:

  • 模型、数据集、训练步数/预算、两臂各自的完整启动参数、随机种子 —— 现在完全看不出这个实验是在什么配置下跑的;
  • KL 曲线和训练稳定性指标(grad-norm / clip-frac / entropy)—— 题目验收明确要求「报告 reward、length、KL 和训练稳定性」,目前只覆盖了 reward 和 length。若本次实验 --kl-coef 0,请说明「未启用 KL 惩罚,此处仅记录 KL 距离」;
  • Length 曲线请按 Correct / Incorrect 分开呈现 —— 你已经把这个指标实现出来了,报告里用上就行。「错误回答不再变长」是 Dr.GRPO 的核心论点(论文 Fig.1),合并成一条总长度曲线看不出这个效果,而这恰恰是最有说服力的证据;
  • 顺带修一下笔误「AIME24qu'xian曲线」。

三、PR 正文没跟着代码更新(必改)

  • 「变更」表仍列着已删除的 examples/algorithms/dr_grpo/run-qwen25-3b-dr-grpo-4xgpu.sh;
  • 「变更」和「已知限制」两处仍写着「在静态 CP 下拒绝原有的 seq-mean-token-mean」,但这段代码已经移除;
  • 正文的 commit 号停在 6bb0798,现在是 10 个 commit;
  • 检查清单里「已完成 CP=1、CP>1、parameter-delta 等价性和短程 paired-run GPU 验收」仍是未勾选 [ ],而上方验证表写「通过」。这两处从上一轮就不一致,请统一。

四、上轮第 6 条的两个测试还没补

  • padding 维度仍零覆盖:max_seq_lens / padded_total_lengths 在 get_sequence_loss_aggregator 里透传了,但所有测试仍传 None。补一条「同一批数据,给不给 padding,loss 逐位相等」。
  • 长短差距仍是 response_lengths=[2, 3](1.5×)。请拉到 [8, 512] 这个量级,并断言 seq-mean-token-mean 与 seq-mean-token-sum-norm 给出不同的长短样本相对梯度权重。目前的测试只验证了新模式自身自洽,没有验证它确实改变了长度加权 —— 而这是整个算法的意义所在。

# ⭐ Feature

## Add paired Qwen3.5 recipe

- Add a Qwen3.5-4B Dr.GRPO recipe with USE_DRGRPO switching.
- Remove the obsolete Qwen2.5-3B model configuration.

---

# ✅ Tests

## Validate fixed-sum normalization

- Cover padding-preserving loss and gradient behavior.
- Assert distinct short/long response weighting at lengths 8 and 512.
Copilot AI review requested due to automatic review settings August 6, 2026 14:29

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@ZiyiTsang

ZiyiTsang commented Aug 10, 2026 •

Copy link
Copy Markdown
Contributor Author

Dr.GRPO Review Response (Second)

感谢 reviewer 对本次 Dr.GRPO 变更进行细致、具体的审阅,下面逐项回复本轮剩余意见。

一、Recipe 与实验配置

Reviewer 意见 回复与处理 对应位置
Dr.GRPO 脚本被删除后没有可运行 recipe。 已补充一个可直接运行的 Qwen3.5-4B Dr.GRPO/GRPO paired recipe。脚本沿用现有 recipe 的启动结构,包含 scripts/entrypoint/local.sh、模型配置、MODEL_DIR / DATA_DIR 环境变量和 Copyright header;通过 USE_DRGRPO 在同一份脚本中切换两臂,共享模型、数据、预算和优化器设置,减少配置漂移。 examples/algorithms/dr_grpo/run-qwen35-4B-dr-grpo-2xgpu.sh
Qwen2.5-3B 配置与上一轮“使用 Qwen3.5 系列”的意见相反。 已移除 Qwen2.5-3B 的正式模型配置,并补充了 Qwen3.5 的模型配置。 scripts/models/qwen35-4B.sh
CISPO recipe 的修改与本题无关。 已将 CISPO 的无关修改排除出本次 Dr.GRPO 交付范围;CISPO recipe 不再作为本次实验复现入口。 examples/algorithms/cispo/
报告中的数学任务实验缺少可追溯的配置。 已在实验说明中补充模型、训练/评测数据、训练预算、随机种子、并行配置以及两臂的完整启动参数,并明确 paired 对照只切换 Dr.GRPO 所需的显式参数组合。 PR/实验说明中的“端到端测试结果”部分

二、对比报告元数据与指标

Reviewer 意见 回复与处理
端到端结果缺少模型、数据集、预算、完整启动参数和随机种子。 已补充完整实验 metadata,包括模型、训练集、AIME24 评测集、训练步数/rollout 预算、batch 配置、优化器、并行策略、随机种子,以及 GRPO 与 Dr.GRPO 两臂各自的完整参数。
缺少 KL 和训练稳定性指标。 已补充 KL 曲线以及 grad-norm、clip-frac、entropy 稳定性指标。实验使用 --kl-coef 0 时,报告明确说明未启用 KL 惩罚;对应曲线只记录 policy 与 reference 之间的 KL 距离,不将其解释为 KL regularization loss。
Length 曲线没有区分正确与错误回答。 我的原本实现本来就会区分正确与错误回答。已经接入 rollout metrics 的分组统计,分别展示 Correct response length 和 Incorrect response length,并在报告中说明该拆分用于直接观察 Dr.GRPO 关注的错误回答长度行为。
报告中有 “AIME24qu'xian曲线” 笔误。 已修正为 “AIME24 曲线”。

三、PR 正文与验证状态

Reviewer 意见 回复与处理
“变更”表仍列出已删除的 run-qwen25-3b-dr-grpo-4xgpu.sh。 已从变更表和相关说明中删除该路径,改为当前 Qwen3.5-4B recipe。
“变更”和“已知限制”仍描述静态 CP 会拒绝原有 seq-mean-token-mean。 已删除这段过时表述。静态 CP 的实现保留原有 seq-mean-token-mean 行为,并通过 step-global token normalizer 支持新的固定尺度路径。
PR 正文的 commit 号停留在 6bb0798。 已将正文更新到当前 10-commit 版本,不再引用旧的起始 commit 作为当前验证版本。
验证表和检查清单对 GPU 验收状态不一致。 已统一两处状态:CP=1、CP>1、parameter-delta 等价性和短程 paired-run GPU 验收在验证表与检查清单中使用同一状态和措辞。

四、补充上轮第 6 条测试

Reviewer 意见 回复与处理 测试覆盖
max_seq_lens / padded_total_lengths 仍为零覆盖。 已新增同一批数据在不提供 padding 参数、提供 padded_total_lengths、以及提供 max_seq_lens 的情况下的对照;测试逐位比较 loss 和梯度,确认 padding 只影响布局,不改变固定尺度目标。 test_padding_kwargs_preserve_fixed_sum_result
[2, 3] 的长度差距不足以验证长度加权变化。 已将测试改为 response_lengths=[8, 512],并同时检查两种 aggregation 的 loss、各 response 的梯度和长样本在总梯度中的相对份额。测试确认 seq-mean-token-mean 对长短 response 等权,而 seq-mean-token-sum-norm 按有效 token 数改变长短样本的相对梯度权重。 test_sum_norm_reweights_short_vs_long_responses

再次感谢 reviewer 指出这些问题。这些修改使 recipe、实验说明、PR 正文和测试之间的对应关系更加清晰,也让固定尺度目标的长度加权行为有了直接的可执行验证。

在此基础上,我还新增了中英doc文档,以便减轻用户认知成本。

待讨论:Dr.GRPO 的接口形态

我们希望与 maintainer 进一步讨论一个接口设计问题:是否应该把 Dr.GRPO 注册为新的 advantage-estimator。我们的倾向是不新增 --advantage-estimator dr_grpo,而是参考 verl 及其他 RL 框架,把 Dr.GRPO 视为可以注入现有实现的独立 trick/目标变体:

讨论点 倾向方案 原因
estimator 命名 保留现有 --advantage-estimator grpo,不额外注册 dr_grpo Dr.GRPO 的核心变化是 advantage 的标准差归一化和 policy-gradient loss 的长度归一化,不一定需要复制一套 estimator 实现。
能力组合 将固定尺度 loss、关闭 group-std normalization 等行为作为可组合的注入项 便于把该 trick 注入任意已有 estimator/目标实现,避免为每个算法增加重复分支。
用户迁移成本 通过配置或 recipe 组合启用,不要求其他用户修改框架核心接口 用户可以在现有训练实现上按需使用,不会因为新增一种算法名称而同步改动 dispatch、注册表或自定义 estimator。
适用范围 具体哪些 estimator 可以安全组合,及参数命名/归属方式,留待 maintainer 讨论后确定 不同 estimator 的数学契约可能不同,需要在保持组合性的同时明确边界,避免产生表面兼容但语义不成立的组合。

当前实现先保留显式、可测试的参数组合,以便验证目标函数和 CP 行为;advantage-estimator 是否维持为基础实现、Dr.GRPO 是否作为通用 trick 注入,则作为下一步 API 设计讨论项。敬请 reviewer 和 maintainer 指正。

# 🐛 Bug Fix

## Make the Dr.GRPO launcher portable

- Remove hard-coded Megatron and virtual-environment paths.
- Resolve model and data roots from `EXP_DIR` or caller-provided environment variables.
- Preserve the local entrypoint, model configuration, and paired GRPO/Dr.GRPO launch options.
# 📝 Documentation

## Document Dr.GRPO in algorithm references

- Explain centered group advantages and fixed-scale token-sum aggregation.
- Document the CLI parameters and Qwen3.5-4B paired recipe.

- Keep English and Chinese algorithm references structurally aligned.
# 🐛 Bug Fix

## Clean the Dr.GRPO launcher

- Remove recipe-local NCCL, FlashInfer, and W&B environment initialization.
- Follow the existing local entrypoint and path configuration pattern.
- Keep the Qwen3.5-4B GRPO/Dr.GRPO training arguments and environment overrides.
# ⭐ Feature

## Integrate upstream algorithm updates

- Add REINFORCE++ implementation, recipes, tests, and bilingual documentation.
- Preserve Dr.GRPO loss aggregation and argument validation while resolving merge conflicts.
- Bring upstream training, logging, distributed utilities, and XPU documentation changes into the branch.
# 🎨 Style

## Align W&B import formatting

- Remove the extra blank line in the W&B adapter imports.
- Place the optional wandb import with the third-party imports in the metrics service.
Copilot AI review requested due to automatic review settings August 10, 2026 09:41

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

# 🐛 Bug Fix

## Align W&B imports with CI formatting

- Restore the blank line separating third-party and first-party imports.
- Keep wandb grouped with FastAPI, Pydantic, and Ray imports.

---

# ✅ Tests

## Validate repository hooks

- Run the full pre-commit suite with generated wandb logs isolated from module discovery.
Copilot AI review requested due to automatic review settings August 10, 2026 14:04

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@li126com

Copy link
Copy Markdown
Member
  1. [P0] fixed-budget 梯度被错误放大 M × DP × CP

seq-mean-token-sum-norm reducer 已经对每个 micro-batch/rank 计算:

  local policy token-loss sum / B

其中 B 是固定 response budget。

但 get_per_token_loss_scale 又返回:

  M / G * (DP × CP) * T

其中:

  • M = num_microbatches
  • G = global batch size
  • T = optimizer step 的全局有效 response token 数

当前 Megatron calculate_per_token_loss 路径的真实行为是:

  • pipeline schedule 在 per-token 模式下不除 M;
  • DDP gradient_scaling_factor=1,跨 DP×CP 使用 SUM,不做 average;
  • finalize_model_grads 最后只除全局 T。

因此最终梯度为:

  M × DP × CP / (G × B) * grad(global policy token-loss sum)

目标应为:

  1 / (G × B) * grad(global policy token-loss sum)

即当前参数更新相对目标放大 M × DP × CP,会随 micro-batch、DP 和 CP 拓扑改变,直接违反
“DP/CP、micro-batch 不改变最终统计”的验收要求。

现有 test_per_token_finalizer_scale_recovers_fixed_dr_grpo_denominator 在 oracle 中手工加入了
真实 per-token schedule 不存在的 /M;CP 测试也手工加入了不存在的 /CP,因此测试通过不能
反证该问题。

建议按 optimizer window 统计实际全局 (N,T):

  • 若 PG reducer 已经除 B,则 compensation 使用 T/N;
  • 或保留原始 token-sum numerator,统一使用 T/(N*B)。

请补充真实 Megatron backward/parameter-delta 回归,覆盖不同 M、DP 和 CP。

涉及位置:

  • relax/backends/megatron/cp_utils.py:get_per_token_loss_scale
  • relax/backends/megatron/loss.py 的 per-token fixed-sum 分支
  • tests/backends/megatron/test_grpo_loss_normalization.py
  1. [P1] 当前仍未作为独立 Dr.GRPO 算法接入

当前 recipe 使用:

  --advantage-estimator grpo
  --disable-grpo-std-normalization
  --pg-loss-aggregation seq-mean-token-sum-norm

registry 中没有 dr_grpo,缺少 --disable-grpo-std-normalization 时也只 warning,仍可运行只实现
一部分 Dr.GRPO 语义的配置。

题目明确要求“作为独立变体接入并与标准 GRPO 清晰区分”。内部实现可以继续复用 GRPO,但
外部应提供原子的:

  --advantage-estimator dr_grpo

并由该入口统一保证 reward centering、no group-std 和 fixed-budget reduction,避免漏配后仍
以 Dr.GRPO 名义运行。

  1. [P1] 当前 PR 正文被另一项 SDPO 任务误覆盖

当前 diff 中没有任何 SDPO、EnvironmentFeedback、opd_sample_mask 或 teacher-update 实现,
但 PR 正文整篇描述的是 Relax-SDPO。

GitHub 正文编辑历史显示:

  • 2026-08-10 06:00:33 UTC:正文仍是 Dr.GRPO;
  • 2026-08-10 09:15:16 UTC:正文被替换为当前 SDPO 内容。

请恢复 Dr.GRPO 正文。需要注意,06:00 版本中的 Qwen3.5-4B reward、length、KL 和稳定性
结果仍标记为“待补图”,因此恢复后还要补充当前 head 对应的完整实验结果、日志和复现信息。

当前状态不满足“设计文档、recipe、标准 GRPO 对比报告”的交付要求。

  1. [P1] train/loss 和 train/pg_loss 没有使用实际反传 scale

policy_loss_function 先保存 reported_loss;外层 get_per_token_loss_scale 只乘到反传 loss,
没有同步乘到 reported loss。随后 model.py 又按全局 token 数 T 对日志做除法。

因此当前:

  train/pg_loss = sum(PG) / (B*T)

它既不是目标 sum(PG)/(N*B),也不代表当前实际施加到参数上的梯度尺度,而且会随有效 token
数变化。使用这些曲线判断 reduction invariance 或训练稳定性会产生误导。

请让 objective-facing 的 loss/pg_loss 日志使用与反传一致的 window scale,同时将
ppo_kl、clipfrac 等纯诊断指标保持为统一、可比较的统计口径。

  1. [P2] fixed aggregation 与辅助 loss 的组合语义不一致

固定 B reducer 只应用于 pg_loss,但外层 compensation 乘到了包含 entropy、explicit KL 和
OPD 的整个 total loss。当前 recipe 将 entropy/KL 系数设为 0,因此不影响这次 paired run;
公共 CLI 却没有禁止这些组合。

请明确契约并二选一:

  • PG、entropy、explicit KL 等统一使用 fixed N*B reduction;
  • 或只对 PG 做 compensation,其他项保持各自定义,并拒绝没有明确定义的组合。

需要为非零 entropy/KL 至少补一个回归测试。

  1. [P3] CISPO recipe 移动与本题无关

当前 diff 仍将 CISPO recipe 移入 examples/algorithms/cispo/,并修改 Usage、source 和 EXP_DIR
相对路径。该脚本仍使用 advantage-estimator cispo,不包含任何 Dr.GRPO 配置,因此不是本题
实现所必需的改动;这也与之前回复中“已排除 CISPO 无关修改”不一致。

建议拆到独立 PR。multimodal training script 仅删除一处尾随空白,可以直接撤掉,但不作为
功能性 review finding。

结论:当前 reducer 的 token-sum/B 局部数学是正确的,主要 correctness blocker 是它接入
Megatron per-token finalizer 时又错误加入了 M 和 DP×CP 补偿。修正该缩放、补真实
parameter-delta 测试、注册独立 estimator,并恢复正确的 PR 正文和实验报告后再验收。

@li126com

li126com commented Aug 17, 2026 •

Copy link
Copy Markdown
Member

感谢作者为 Dr.GRPO 做的实现和多轮修改。这个 PR 在中英文公式文档、长短 response 混合测试、padding/CP 覆盖、Qwen3.5 recipe,以及 Correct/Incorrect response length 指标等方面投入了很多工作;作者也认真回应并修正了多项 review 意见,这些内容对后续 Dr.GRPO 实现很有参考价值。

综合当前代码和题目验收标准,本轮验收不通过,主要原因如下。

1. Fixed-budget loss 的梯度缩放仍然不正确

当前实现先将 policy loss 按固定 response budget B 归一化,之后又在外层加入了包含 micro-batch 数量、DP×CP world size 和 token 数 T 的补偿系数。

但 Relax 当前使用的 Megatron per-token 训练路径中:

  • pipeline schedule 不会再除以 micro-batch 数量;
  • DDP 对 DP×CP rank 的梯度执行 SUM;
  • finalizer 最后只除以全局 token 数 T。

因此当前最终梯度会额外放大 num_microbatches × DP × CP,训练结果会随着 micro-batch 划分和 DP/CP 拓扑变化,不再是目标公式中的 Σloss / (N × B)。这是 Dr.GRPO 核心 loss aggregation 的训练正确性问题,也是本次不通过验收的主要原因。

现有测试通过手工除以 micro-batch 或 CP world size 模拟了实际 Megatron 中不存在的 averaging,因此测试虽然通过,但没有覆盖真实 schedule、DDP 和 finalizer 组合后的参数梯度。这里需要删除多余的缩放,并补充真实 gradient 或 parameter-delta 对照测试。

2. 尚未作为独立 Dr.GRPO 变体接入

题目要求 Dr.GRPO 作为独立变体接入,并与标准 GRPO 清晰区分。当前实现仍然使用:

--advantage-estimator grpo

再组合关闭 reward std normalization、修改 PG loss aggregation 等配置来启用 Dr.GRPO,registry.py 中也没有独立的 dr_grpo estimator/algorithm entry。

这会允许用户只打开部分配置,形成既不是标准 GRPO、也不是完整 Dr.GRPO 的中间状态。当前参数校验部分场景也只是 warning,没有从配置层保证 Dr.GRPO 的 advantage 和 loss 契约同时生效。因此这一项与“独立变体接入”的验收要求尚未对齐。

3. 可复现对比报告还没有完整对应当前实现

PR 已经补充了一些实验说明和 recipe,这是很好的方向,但当前证据仍未形成一组可从当前提交复现的标准 GRPO 与 Dr.GRPO 对比:

  • 报告中的实验使用历史 Qwen2.5-3B 配置,而当前提交的 recipe 是 Qwen3.5-4B;
  • 部分模型、数据路径和 seed 信息仍不完整;
  • PR 描述记录的代码 commit、CP 配置与当前 head/recipe 不一致;
  • GPU parameter-delta 验证仍标记为 pending;
  • 尚缺少完整对应当前实现的 reward、response length、KL 和训练稳定性结果。

因此目前还不能确认两个算法是在相同模型、数据、随机种子和训练预算下完成的可复现对比。

综上,本次结论是针对 Dr.GRPO 任务的核心公式正确性、独立配置入口和可复现实验三项验收要求,并不否定 PR 中已经完成的文档、测试和工程工作。感谢作者投入时间参与 Relax 社区贡献,期待看到你之后更多高质量的社区贡献。

@SigureMo SigureMo closed this Sep 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants