Conversation
# ⭐ 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
# ✅ 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.
|
评审意见:Request changes 数学实现是对的,问题集中在集成影响半径和报告与代码对不上。
|
# ♻️ 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
对 review comment 的处理
|
# ⭐ 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.
|
实验结果已就绪 |
There was a problem hiding this comment.
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 raiseTypeErrorwhenSample.rewardis a dict and--reward-keyis unset (becauseSample.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.
| import logging | ||
| import math | ||
| from numbers import Real | ||
| from typing import Any, Literal | ||
|
|
| # 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, | ||
| } |
| # 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 |
There was a problem hiding this comment.
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/Relaxstep. 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-specificsource/exportcommands 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 raisesTypeErrorwhen a sample reward isNoneor a dict (common when--reward-keyisn’t set). Since this function is called unconditionally fromcompute_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."
There was a problem hiding this comment.
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 useget_logger(__name__)rather thanlogging.getLogger. This file currently has neither (it importsloggingdirectly and initializesloggervialogging.getLogger).
import logging
import math
from numbers import Real
examples/algorithms/cispo/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh:31
EXP_DIRdefault path is one directory too deep after moving this script underexamples/algorithms/cispo/. With the current../../../../../exps, the default points to the grandparent of the repo (likely unintended) instead of matching the../../../../expspattern 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 patchesrelax.utils.metrics.adapters.wandb.wandb.init. However,MetricsService._init_wandb()currently callswandb.init()fromrelax.utils.metrics.servicedirectly and does not pass these kwargs, so the patch/expectations won't match and the test will fail. Either updateMetricsService._init_wandb()to reuse the existinginit_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"
| 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 |
# 🔩 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.
# 🐛 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.
|
感谢更新,已处理的review意见:
还剩下面几项,其中第一、二项是这道题的核心交付物,请优先处理。 一、recipe 与实验对不上(必改) 现在有四处互相矛盾:
题目的交付物明确包含 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,最好在同一个脚本里用一个开关切换,避免配置漂移。 二、对比报告缺元数据(必改) 「端到端测试结果」目前只有一句话加三张图。请补上:
三、PR 正文没跟着代码更新(必改)
四、上轮第 6 条的两个测试还没补
|
# ⭐ 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.
Dr.GRPO Review Response (Second)感谢 reviewer 对本次 Dr.GRPO 变更进行细致、具体的审阅,下面逐项回复本轮剩余意见。 一、Recipe 与实验配置
二、对比报告元数据与指标
三、PR 正文与验证状态
四、补充上轮第 6 条测试
再次感谢 reviewer 指出这些问题。这些修改使 recipe、实验说明、PR 正文和测试之间的对应关系更加清晰,也让固定尺度目标的长度加权行为有了直接的可执行验证。 在此基础上,我还新增了中英doc文档,以便减轻用户认知成本。 待讨论:Dr.GRPO 的接口形态我们希望与 maintainer 进一步讨论一个接口设计问题:是否应该把 Dr.GRPO 注册为新的
当前实现先保留显式、可测试的参数组合,以便验证目标函数和 CP 行为; |
# 🐛 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.
# 🐛 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.
seq-mean-token-sum-norm reducer 已经对每个 micro-batch/rank 计算: 其中 B 是固定 response budget。 但 get_per_token_loss_scale 又返回: 其中:
当前 Megatron calculate_per_token_loss 路径的真实行为是:
因此最终梯度为: 目标应为: 即当前参数更新相对目标放大 M × DP × CP,会随 micro-batch、DP 和 CP 拓扑改变,直接违反 现有 test_per_token_finalizer_scale_recovers_fixed_dr_grpo_denominator 在 oracle 中手工加入了 建议按 optimizer window 统计实际全局 (N,T):
请补充真实 Megatron backward/parameter-delta 回归,覆盖不同 M、DP 和 CP。 涉及位置:
当前 recipe 使用: registry 中没有 dr_grpo,缺少 --disable-grpo-std-normalization 时也只 warning,仍可运行只实现 题目明确要求“作为独立变体接入并与标准 GRPO 清晰区分”。内部实现可以继续复用 GRPO,但 并由该入口统一保证 reward centering、no group-std 和 fixed-budget reduction,避免漏配后仍
当前 diff 中没有任何 SDPO、EnvironmentFeedback、opd_sample_mask 或 teacher-update 实现, GitHub 正文编辑历史显示:
请恢复 Dr.GRPO 正文。需要注意,06:00 版本中的 Qwen3.5-4B reward、length、KL 和稳定性 当前状态不满足“设计文档、recipe、标准 GRPO 对比报告”的交付要求。
policy_loss_function 先保存 reported_loss;外层 get_per_token_loss_scale 只乘到反传 loss, 因此当前: 它既不是目标 sum(PG)/(N*B),也不代表当前实际施加到参数上的梯度尺度,而且会随有效 token 请让 objective-facing 的 loss/pg_loss 日志使用与反传一致的 window scale,同时将
固定 B reducer 只应用于 pg_loss,但外层 compensation 乘到了包含 entropy、explicit KL 和 请明确契约并二选一:
需要为非零 entropy/KL 至少补一个回归测试。
当前 diff 仍将 CISPO recipe 移入 examples/algorithms/cispo/,并修改 Usage、source 和 EXP_DIR 建议拆到独立 PR。multimodal training script 仅删除一处尾随空白,可以直接撤掉,但不作为 结论:当前 reducer 的 token-sum/B 局部数学是正确的,主要 correctness blocker 是它接入 |
|
感谢作者为 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 训练路径中:
因此当前最终梯度会额外放大 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 对比:
因此目前还不能确认两个算法是在相同模型、数据、随机种子和训练预算下完成的可复现对比。 综上,本次结论是针对 Dr.GRPO 任务的核心公式正确性、独立配置入口和可复现实验三项验收要求,并不否定 PR 中已经完成的文档、测试和工程工作。感谢作者投入时间参与 Relax 社区贡献,期待看到你之后更多高质量的社区贡献。 |
总结
本PR为 Relax 增加了可显式组合的 Dr.GRPO 目标,复现论文
Dr.GRPO 中的两个核心修改:
对于包含 (G) 个回复的 prompt group,Dr.GRPO 使用中心化 advantage:
而不是标准 GRPO 中带组内标准差的 advantage。实现的目标为:
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变更
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.pyrelax/backends/megatron/data.pyrelax/backends/megatron/loss.pytests/backends/megatron/test_grpo_loss_normalization.pytests/utils/test_arguments_dr_grpo.pyscripts/models/qwen25-3B.shexamples/algorithms/dr_grpo/run-qwen35-4B-dr-grpo-2xgpu.shCLI、配置与兼容性变化
新增显式参数:
Dr.GRPO 由以下参数显式组合:
--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
e8c3169(当前 HEAD;完整值以git rev-parse HEAD为准)可复制命令
聚焦单元测试:
语法与静态检查:
新增测试
test_response_length_normalization_preserves_existing_behaviortest_seq_mean_token_sum_norm_uses_one_scale_factor_for_all_responsestest_seq_mean_token_sum_norm_requires_positive_scale_factortest_per_token_finalizer_scale_recovers_fixed_dr_grpo_denominatortest_per_token_finalizer_cp_shards_recover_fixed_dr_grpo_denominatortest_per_token_finalizer_requires_step_global_not_microbatch_normalizertest_real_megatron_static_iterator_reuses_step_normalizertest_real_megatron_dynamic_iterator_reuses_step_normalizertest_static_cp_dr_grpo_matches_cp_one_fixed_scale_gradienttest_padding_kwargs_preserve_fixed_sum_resultmax_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_explicittest_pg_loss_aggregation_rejects_invalid_combinationstest_explicit_pg_loss_scale_factor_overrides_rollout_max_response_length--pg-loss-scale-factor优先于--rollout-max-response-len。单元、集成测试结果
7 passed, 12 skipped;跳过项依赖当前开发环境未提供的 Megatron/CPU process-group 条件。git diff --checkpre-commit run --all-files --show-diff-on-failurepre-commit,未执行。端到端测试结果
端到端结果分为两个实验。每个实验都必须同时记录 reward、Correct/Incorrect length、KL 和训练稳定性指标。
本次两个实验均使用
--kl-coef 0,未启用 KL 惩罚;报告中的 KL 曲线仅记录 KL 距离。/data/share/Qwen2.5-3B(历史配置;不作为当前 recipe)xxxxxxx/math_deepmath_deal.jsonlxxxxxxxx/aime24/test.jsonl--num-rollout 200,--global-batch-size 512,每 prompt 8 samples1e-6,constant schedule,weight decay 0,clip grad 1.0--disable-grpo-std-normalization、--pg-loss-aggregation seq-mean-token-sum-norm两个参数上不同AIME24 曲线:
Reward 曲线:
Length 曲线:
风险与回退
已知限制
seq-mean-token-sum-norm要求--advantage-estimator grpo。--fully-async,因为 streaming iterator 无法安全建立optimizer-step-global token normalizer。
--custom-pg-loss-reducer-function-path组合使用。不能将二者当作数值上等价的实现。
风险
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行为。检查清单
docs/刻意未修改:本文件为 PR 材料,不是发布文档。More info: issue #86