Skip to content

【No.30】feat(dr-grpo): add Dr.GRPO support - #239

Open
idvzchusvol wants to merge 10 commits into
redai-infra:mainfrom
idvzchusvol:feat/dr-grpo-main
Open

【No.30】feat(dr-grpo): add Dr.GRPO support#239
idvzchusvol wants to merge 10 commits into
redai-infra:mainfrom
idvzchusvol:feat/dr-grpo-main

Conversation

@idvzchusvol

@idvzchusvol idvzchusvol commented Aug 4, 2026

Copy link
Copy Markdown

feat(dr-grpo): add Dr.GRPO support

What

本 PR 为 Relax 的 Megatron 训练后端增加 Dr.GRPO 支持。

主要改动包括:

  • 算法注册与服务拓扑:在算法注册表和 --advantage-estimator 中加入 dr_grpo,并复用 GRPO 现有的 Rollout、Actor、Advantages、Reference 和 ActorFwd 服务拓扑。
  • Reward 与 advantage:实现 Dr.GRPO 的组内 reward 中心化语义,即减去组内 reward 均值,但不除以组内标准差。
  • Actor loss reduction:使用固定 response budget 归一化 Actor loss:

$$L_{\mathrm{Dr.GRPO}} = \frac{S_{\mathrm{actor}}}{N B}$$

其中 $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 语义。

  • 有效 token 计数:统一 cp_size == 1cp_size > 1 的有效 token 定义,并处理全局有效 token 数为 0 的 optimizer window。
  • 文档与测试:补充中英文使用文档、Dr.GRPO 单元测试、Qwen3.5-4B paired recipe。

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 使用:

$$A_i = r_i - \frac{1}{G}\sum_{j=1}^{G} r_j$$

与 GRPO 不同,Dr.GRPO 不再使用组内 reward 标准差除以 $A_i$

去掉方差归一化是 Dr.GRPO 的核心,因此 --normalize-advantages 会被参数校验拒绝:该选项会在组内中心化之后再叠加一层全局 whitening,与之矛盾。

Fixed-budget loss reduction

这里的 optimizer window 是一次 optimizer.step() 对应的完整梯度累积范围,包含该次参数更新前参与 forward/backward 的所有 micro-batch。下文的 $N$$T$ 都在这个范围内统计。

令:

  • $S$ 为 optimizer window 内组合后的 Actor token loss 分子;
  • $T$ 为全局有效 response token 数;
  • $N$ 为全局 response 数;
  • $B$rollout_max_response_len

Relax 在 optimizer window 开始前计算:

$$\alpha = \frac{T}{N B}$$

随后将组合后的 Actor loss 乘以 $\alpha$,并复用 Megatron 原有的全局 token normalization:

$$\frac{\alpha S}{T}=\frac{T}{N B}\frac{S}{T}=\frac{S}{N B}$$

这样既能满足 Dr.GRPO 的固定预算目标,也能保留 Megatron 在 CP 场景下要求的 calculate_per_token_loss 路径。

具体的实现逻辑集中在 optimizer-window metadata provider 中:

  • prepare_policy_optimizer_window_metadata 在每次参数更新前统计全局 $(N,T)$,计算 $T/(N B)$,并生成专用的 __dr_grpo_window_scale__
  • MegatronTrainRayActor 将该 scale 绑定到当前 optimizer window;bind_optimizer_window_metadataDataIterator.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 的分母与 $\alpha$ 中的
    $T$ 保持一致。
  • 对于全局 T = 0 的 optimizer window,将跳过 optimizer 和 scheduler,同时报告 warning 信息。
  • metric denominator 在 T = 0 时将报告 0,不再出现除零错误;被跳过的窗口 grad_norm 保持 0.0 而非 NaN,避免被其他检查误判。

Testing

单元测试

Dr.GRPO 测试覆盖:

  • 算法注册和参数选择。
  • Dr.GRPO 的组内 reward 仅做中心化,不再除以组内标准差。
  • Megatron 与 Advantages service 两条路径均不添加 Dr.GRPO 独有的 reward-side KL。
  • 独立 Dr.GRPO scale key 与 fully-async streaming scale 的隔离。
  • Optimizer-window 全局 $(N,T)$ 统计及 current-device reduction。
  • 真实分布式 process group 下的 DP/CP metadata reduction。
  • CP padding、loss mask 和长短 response 的相对 token 权重。
  • Fixed micro-batch metadata replay 和 rebind offset 生命周期。
  • Response budget、fully-async、reward centering 和 KL 参数校验。
  • GRPO 默认路径回归。
  • Mixed-length/topology reduction oracle。
  • 全 mask response 的有效 token 计数在 cp_size = 1/2/4 下一致为 0。
  • 空 optimizer window 的标记、投递到每个 micro-batch,以及整窗全 mask 时计数为 0。
  • --normalize-advantages 与 Dr.GRPO 组合时被拒绝。
  • 生产 train_one_step 在空窗口下不调用 optimizer 与 scheduler 的 .step(),且 grad_norm0.0 而非 NaN;正常窗口下两者各调用一次。
  • Megatron finalize_model_grads1/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 负载不均衡。具体而言:

N = 16, B = 1024, T = 4930, S = 402
alpha    = T / (N * B) = 0.3009033203125
Dr.GRPO  = S / (N * B) = 0.0245361328125
GRPO     = S_grpo / T  = 0.031580906737842154

而六种 (DP, CP) 拓扑 × 两种 micro-batch size 下,Dr.GRPO 的实测 loss 与上表完全一致。

同一份窗口还在每个拓扑上顺序覆盖:

  • Dynamic batching:不规则的 micro-batch 分组不改变 window scale。
  • 显式 paddingtp * cp * 2 对齐布局与普通 THD 隐式 padding
    命中相同的 T 与 loss,padding token 不进入任何分子或分母。
  • Entropy 与 explicit KL:两者与 policy gradient 使用同一 fixed-budget
    reduction,组合公式成立。
  • OPSM:masking 后分母仍为完整的 T,clip fraction 与 PG 分子命中冻结值。
  • 长短 response 的梯度权重:梯度正比于有效 token 数,覆盖约 73x 的长度跨度。
  • 单条全 mask response:被完全 mask 的 response 计 0 个 token,finalizer 分母与 T
    保持一致,finalizer 之后的梯度不随 CP degree 变化。

并冻结以下 contract:

  • (N, T) 只在 DP-without-CP group 上归约,不重复统计 CP replica。
  • GRPO 路径不携带 Dr.GRPO 的 window scale。
  • 仅改变 configured rollout_max_response_len 时 Dr.GRPO 结果按比例变化、GRPO 不变,
    以区分 configured budget 与 observed max response length。
  • 梯度冻结绝对值而非仅冻结比值,覆盖 backward 与 reported metric 两条独立的 scale 应用路径。

测试结果

pytest --collect-only -q \
  tests/backends/megatron/test_dr_grpo.py \
  tests/utils/test_arguments_dr_grpo.py
41 tests collected in 23.13s

pytest -q \
  tests/backends/megatron/test_dr_grpo.py \
  tests/utils/test_arguments_dr_grpo.py
41 passed, 26 warnings in 323.52s (0:05:27)

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:

Algorithm Reward mean / last20 Accuracy Correct length mean / last20 pooled Incorrect length mean / last20 pooled
GRPO 0.8375 / 0.9250 91.88% 1179.7 / 1292.2 2329.6 / 2565.9
Dr.GRPO 0.8300 / 0.9437 91.50% 733.6 / 619.6 2547.4 / 1173.2

Gradient:

Algorithm Grad norm mean / p95 / max
GRPO 0.2390 / 1.3441 / 2.2343
Dr.GRPO 0.1225 / 0.6507 / 1.9966
dr_grpo_qwen35_4b_gsm8k_step0_199

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/10 steps,指标全部 finite,无 OOM、collective hang 或 Ray job failure。全零 mask 的窗口按预期跳过 optimizer 与 scheduler,grad_norm0.0 而非 NaN;半置零窗口不触发该分支,optimizer 正常执行。

  • pre-commit run --all-files passes
  • Tests pass (pytest tests/)
  • New tests added (if applicable)
  • Documentation updated (if applicable)

Type of Change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation update
  • Refactoring (no functional changes)
  • Performance improvement
  • CI/CD or build changes

Copilot AI lite review requested due to automatic review settings August 4, 2026 11:52

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 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_grpo and 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.

Comment on lines +705 to +707
"""Prepare optimizer-window metadata consumed by the policy loss."""
if args.advantage_estimator == "dr_grpo":
step_stats = []

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

已在 69e1dbd 修复。

  • _validate_dr_grpo_args() 在启动时校验 rollout_max_response_len 必须是正整数;
  • prepare_policy_optimizer_window_metadata() 在执行 T / (N * B) 的消费边界保留同样的防御性校验;
  • tests/utils/test_arguments_dr_grpo.py 覆盖 0True,provider 测试覆盖 zero budget。

这样 CLI 与程序化/helper 调用都会 fail fast,不会再延迟到除法处报 TypeErrorZeroDivisionError

Comment on lines +728 to +730
for iterator in data_iterators:
iterator.metadata_by_microbatch = metadata_by_microbatch

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

已在 69e1dbd 修复。

bind_optimizer_window_metadata() 现在无论绑定新 metadata 还是 clear metadata 都会把 microbatch_offset 重置为 0;get_next() 也只在实际存在 metadata 时递增该 offset。新增 rebind regression test:先消费一个 micro-batch,再绑定新 metadata,确认从新 window 的第一个 entry 开始 replay。

Comment thread relax/backends/megatron/data.py Outdated
self.offset = 0
self.microbatch_offset = 0

def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]:

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

已修正。当前 DataIterator.get_next() 的返回类型是 dict[str, Any],与 batch 中可能包含 list、tensor 和 optimizer-window scalar metadata 的实际 contract 一致。

Copilot AI review requested due to automatic review settings August 4, 2026 12:25

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.

Comment thread relax/backends/megatron/loss.py Outdated
Comment on lines +707 to +726
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)

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

已在 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 端到端路径也已完成。

@idvzchusvol idvzchusvol changed the title feat(dr-grpo): add Dr.GRPO support 【No.30】feat(dr-grpo): add Dr.GRPO support Aug 5, 2026
@li126com

li126com commented Aug 6, 2026

Copy link
Copy Markdown
Member

感谢这份提交,亮点:

  • 真正作为独立变体接入。--advantage-estimator dr_grpo + registry.py 新增条目,与题目「作为独立变体接入并与标准 GRPO 清晰区分」完全对齐。registry 的改动是纯新增,没有动 PPO 或其他算法的角色集。
  • 白名单传播做得很干净。我对照 dev 上 10 处含算法名判断的位置逐一核过,该改的都改了 —— 包括 utils/utils.py:440 那份 load_debug_rollout_data 里的第二份拷贝,这处很容易漏。loss.py 和 components/advantages.py 两条 advantage 路径也都覆盖了。
  • std 归一化的处理很精巧。把 dr_grpo 加进 utils.py:185 的外层中心化白名单、但不加进 :206 的 std 门,靠结构本身保证「只中心化、不除 std」,没有引入新分支。
  • DP/CP/micro-batch/padding 不变性的证据是本题最扎实的。16 条 response(1024/256/64/16 各 4 条)+ 不同 prompt 长度和 mask 的确定性 batch,给出 oracle 402/(16×1024) = 0.0245361328125,然后 CP1/DP1、CP1/DP2、CP1/DP4、CP2/DP1、CP2/DP2、CP4/DP1 六种拓扑加两种 micro_batch_size 绝对误差全为 0,GRPO 回归 max abs err < 1.7e-8。这直接把题目「构造长短 response 混合 batch」和「DP/CP、micro-batch 和 padding 不改变最终统计」两条验收打满了。
  • recipe 完全符合仓库惯例:Copyright header、source scripts/entrypoint/local.sh、source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh"(用仓库既有的最小模型配置,没有另造)、环境变量驱动且用 :? 做必填校验。尤其是用 ADVANTAGE_ESTIMATOR 一个开关切两臂,避免了对照实验的配置漂移 —— 这正是对照实验该有的写法。
  • 中英文文档齐全(docs/{zh,en}/guide/dr-grpo-training.md + vitepress sidebar 接线 + examples/algorithms.md + configuration.md),pre-commit run --all-files 已通过。

reveiw 修改意见:

  1. 必改:一处无守卫的 loss scale 会打到 3 个现有配方

relax/backends/megatron/loss.py 的 per-token 分支:

 else:
     if is_dummy:
         loss = 0.0 * loss
  •    elif explicit_loss_scale is not None:
    
  •        loss = loss * explicit_loss_scale
       # Non-dummy per-token path: do NOT scale by cp_size. ...
    

这里没有 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
scripts/training/multimodal/run-qwen35-9B-8xgpu-openr1mm-async.sh
examples/algorithms/run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh

这些任务今天在 per-token 分支是故意不乘 explicit_loss_scale 的(紧跟着的那段注释解释了原因),改动后会被静默多乘一个 scale。

一个佐证:紧邻的 metrics 缩放你是加了守卫的 ——

  • if args.advantage_estimator == "dr_grpo" and explicit_loss_scale is not None:
  •    for key in ("loss", "pg_loss", "entropy_loss", "kl_loss"):
    

说明你知道这里需要守卫,只是漏在了 loss 那一行。请补上同样的条件,或者给 Dr.GRPO 的 window scale 用一个独立的 key(如 dr_grpo_window_scale),避免和 streaming 的 loss_scale 共用同一通道。

  1. 必改:报告缺 length 和 KL 两项

题目验收明确要求「报告 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 的比值本来就会整体缩放梯度尺度,这不是不稳定,但读者容易误读,最好说明两臂的有效学习率是否可比。

  1. 建议确认:Dr.GRPO 独有的 reward-side KL

if args.advantage_estimator == "dr_grpo" and args.kl_coef != 0:
returns = [ret - args.kl_coef * k.detach() for ret, k in zip(returns, kl, strict=True)]

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 所以没暴露。请说明这是刻意对齐某个参考实现,还是应该移除;如果保留,请在文档和对照实验协议里显式声明。

  1. 建议确认:--disable-rewards-normalization 对 dr_grpo 静默失效
  • if args.advantage_estimator in [..., "dr_grpo", ...] and (
  •    args.rewards_normalization or args.advantage_estimator == "dr_grpo"
    

or args.advantage_estimator == "dr_grpo" 让 Dr.GRPO 的中心化无视 --disable-rewards-normalization。语义上是对的(中心化是 Dr.GRPO 的定义),但用户传了这个 flag 却没有任何提示。建议加一条 warning,或在参数校验阶段直接拒绝这个组合。

  1. Copilot 的四条

四条都是有效的,尤其第四条(设备不匹配)建议优先处理:

(1) prepare_policy_optimizer_window_metadata() 未校验 --rollout-max-response-len 为正 → T/(N·B) 会 ZeroDivision/TypeError。
(2) bind_optimizer_window_metadata() 不重置 microbatch_offset → metadata replay 可能错位。
(3) DataIterator.get_next() 的返回类型注解已过时。
(4) loss_masks 可能是 CPU tensor,而 DP group 是 NCCL-backed,dist.all_reduce(stats) 会有设备/类型不匹配,loss_scale 也可能落在 CPU 而 loss 在 CUDA。你的 4 卡实验能跑通说明实际路径上 mask 在 GPU,但建议按 Copilot 的建议把 stats 显式建在当前 device 上,不要依赖上游 dtype/device 的隐式约定。

  1. 其他
  • mergeable_state: dirty,请先 rebase 解决冲突。
  • 另外最好能补充 train_rollout_prob_abs_diff 指标的对比结果,也就是训推 mismatch 的对比结果。此外优先推荐 qwen35 模型的训练对比。

Copilot AI review requested due to automatic review settings August 10, 2026 03:52

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.

Copilot AI review requested due to automatic review settings August 10, 2026 03:56

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.

idvzchusvol commented Aug 10, 2026

Copy link
Copy Markdown
Author

@li126com 感谢您的 review。已在 69e1dbd 完成代码、测试、recipe 和文档修复。Qwen3.5-4B 对照结果及曲线同步到 PR 正文;

逐项回复如下:

  1. Fully-async __loss_scale__ 回归

    • 采用建议的独立 key 方案:Dr.GRPO 只生成和消费 __dr_grpo_window_scale__;streaming __loss_scale__ 保持原有语义。
    • Per-token 非 Dr.GRPO 路径不会消费 streaming scale。
    • Pure fully-async Dr.GRPO 会在参数校验阶段被拒绝;hybrid 在 optimizer window 闭合并合并后使用 Dr.GRPO metadata。
  2. Length、KL 与稳定性报告

    • PR 正文使用 Qwen3.5-4B、GSM8K、CP1/DP2 和 4 张 H20 的 step 0..200 reporting window。该窗口包含 201 个 train/rollout records、每步 16 条 response,共 3216 responses。测试结果已展示到PR正文中。
  3. Dr.GRPO 独有 reward-side KL

    • 已从 Megatron 和 Advantages service 两条执行路径移除。
    • dr_grpo + kl_coef != 0 现在会在启动时直接拒绝;需要 KL penalty 时使用 explicit KL loss。
  4. --disable-rewards-normalization 静默失效

    • 选择 fail-fast:Dr.GRPO 配置该 flag 时参数校验直接拒绝,并说明 group centering 是算法定义的一部分。
  5. Copilot 四项

    • Response budget:参数入口和 metadata provider 双层正整数校验。
    • Metadata offset:bind、rebind 和 clear 均重置 microbatch_offset,且无 metadata 时不递增。
    • Return annotation:已更新为 dict[str, Any]
    • NCCL device:只将 mask scalar sum 搬到 current device,再使用 DP-without-CP group 归约,确保 Dr.GRPO scale 与 loss 位于同一 device。
  6. Rebase 与 Qwen3.5 recipe

    • 分支已 rebase 到当前 main。
    • 新增 scripts/training/text/run-qwen35-4B-4xgpu-dr-grpo.sh,GRPO/Dr.GRPO 只通过 ADVANTAGE_ESTIMATOR 切换。

Copilot AI review requested due to automatic review settings August 10, 2026 04: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.

Copilot AI review requested due to automatic review settings August 13, 2026 07: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

重新复核后,Dr.GRPO 的主体算法实现没有发现新的 correctness 问题:独立 estimator、reward
只做组内中心化、optimizer-window T/(N*B) scale、DP/CP reduction、padding 和 micro-batch
metadata replay 均是自洽的。

当前剩余问题集中在验收证据与对比报告。

  1. KL 对比使用了不同的归一化口径

Dr.GRPO 路径会将 kl_loss 乘 dr_grpo_window_scale = T/(N*B),日志汇总随后再除 T,
因此 Dr.GRPO 的 train/kl_loss 实际为:

sum(KL) / (N*B)

GRPO 未应用该 scale,得到的是:

sum(KL) / T

所以报告中的 0.1692 vs 0.0423 不能直接解释为 Dr.GRPO 的 policy-reference 平均 KL 更小。
这不影响当前 kl-loss-coef=0 的训练结果,但影响验收报告。请额外记录一个两组实验统一定义的
reference KL,例如全局 sum(KL)/T。

  1. 正文声称的核心 regression 没有对应入库测试

正文报告了 16 条 response(1024/256/64/16)、六种 DP/CP 拓扑、两种 micro-batch size、
固定 oracle 以及 GRPO regression <1.7e-8,但当前测试只覆盖 8/4 和 1/4 token 的小例子;
真实 process-group 测试只验证 metadata all-reduce,没有执行正文所述的最终 loss/gradient
oracle。仓库中也找不到 4930、402、0.0245361328125 或 1.7e-8 对应的数据和断言。

另外,当前两个 focused test 文件静态展开为 21 cases,与正文的 35 passed 不一致。请把
deterministic mixed-length/topology oracle 测试或验证脚本入库,并更新真实测试结果。

  1. 200-step 结果缺少精确复现材料

Recipe 能启动同类实验,但尚未固定模型 checkpoint revision、GSM8K split/转换方式和数据
hash,也没有提交逐步日志、summary artifact 或生成表格/图片的脚本。因此目前无法从 PR
内容重建正文中的 reward、length、truncation、KL 和 grad-norm 数字。请补充这些复现信息。

以上不否定算法实现本身;修复重点是让回归测试和对比报告与当前代码及验收要求一致。

Copilot AI review requested due to automatic review settings August 14, 2026 17: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.

@idvzchusvol

Copy link
Copy Markdown
Author

@li126com 感谢您的 review,三个问题已逐条修复。

1. KL 对比口径

我们已改为按统一定义 sum(KL) / T 对比GRPO和Dr.GRPO,见正文的reference KL比较图。

2. 回归测试

抱歉在上一次提交中忽略了相关的测试,我们现在已把 deterministic mixed-length/topology oracle 入库,合并进
tests/backends/megatron/test_dr_grpo.py,测试数量也进行了更正。

  • 使用一个固定的 mixed-length 窗口:16 条 response、1024/256/64/16 四种长度,配不
    规则的 prompt 长度与多种 loss mask 形态,并让 DP rank 之间的 token 负载不均衡;
  • 六种 (DP, CP) 拓扑 × 两种 micro-batch size,验证最终 loss 不随拓扑与 micro-batch 切分改变;
  • 确认 Dr.GRPO 的 scale 在 loss_function 中反向传播和
    日志指标两条路径上都正确生效;
  • 另外覆盖 dynamic batching 的不规则分组、显式 padding 布局、entropy 与 explicit KL
    的组合,以及 OPSM 强制分支。

3. 200-step 复现材料

新增 docs/public/dr-grpo/,参照 docs/public/reinforce-plus-plus/ 的组织方式:

  • experiment_manifest.json:模型 Qwen/Qwen3.5-4B 及 40 位 revision、数据集
    openai/gsm8k 的 revision / config / split / 转换步骤 / 输出 sha256、并行度、
    全部超参与 seed、依赖版本;
  • evidence_index.csv:两个 run 的 job id、状态、时长,以及训练日志的 sha256 与
    rollout 结果的聚合 sha256(日志与逐 response 结果体积较大,没有提交到PR中);
  • training_metrics_long.csv:逐 step 指标;training_curves.svg:正文图表来源。

按照 Relax 的常见做法,我们暂时没有提交直接生成表格/图片的脚本,而是在 manifest 中
给出了正文数字的复现说明。

@li126com

Copy link
Copy Markdown
Member

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 设为零:

  • relax/utils/utils.py#L113-L130 (

    Relax/relax/utils/utils.py

    Lines 113 to 130 in 4b49a79

    # loss mask
    # TODO: compress the loss mask
    loss_masks = []
    for sample in samples:
    # always instantiate loss_mask if not provided
    if sample.loss_mask is None:
    sample.loss_mask = [1] * sample.response_length
    else:
    # NOTE(jiajia): loss_mask is not None only if args.mask_offpolicy_in_partial_rollout is True, so we need to pad it to response_length.
    sample.loss_mask += [1] * (sample.response_length - len(sample.loss_mask))
    assert len(sample.loss_mask) == sample.response_length, (
    f"loss mask length {len(sample.loss_mask)} != response length {sample.response_length}"
    )
    if sample.remove_sample:
    sample.loss_mask = [0] * sample.response_length
    loss_masks.append(sample.loss_mask)
    train_data["loss_masks"] = loss_masks
    )

Dr.GRPO window metadata 使用真实的 mask.sum() 计算 T:

  • relax/backends/megatron/loss.py#L733-L767 (
    def prepare_policy_optimizer_window_metadata(
    args: Namespace,
    rollout_data: RolloutBatch,
    step_local_sample_counts: list[int],
    ) -> list[dict[str, torch.Tensor]] | None:
    """Prepare optimizer-window metadata consumed by the policy loss."""
    if args.advantage_estimator == "dr_grpo":
    if type(args.rollout_max_response_len) is not int or args.rollout_max_response_len <= 0:
    raise ValueError("--rollout-max-response-len must be a positive integer for Dr.GRPO.")
    stats_device = device_utils.make_current_torch_device()
    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(device=stats_device, 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)
    return [{"__dr_grpo_window_scale__": scale} for scale in step_loss_scales]
    )

但传给 Megatron finalizer 的 token count 定义不一致:

  • CP=1:对每条 response 使用 clamp_min(mask.sum(), 1);
  • CP>1:使用真实、未 clamp 的 CP-local 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)。

最小例子:

  • N=2, B=4
  • masks 为 [1, 1] 和 [0, 0]
  • numerator S=2
  • metadata 得到 T=2,所以 alpha=T/(N*B)=1/4

结果:

  • CP=1 finalizer token count 为 2 + 1 = 3,最终为 (alpha*S)/3 = 1/6
  • CP=2 汇总 token count 为 2,最终为 (alpha*S)/2 = 1/4
  • 正确目标 S/(N*B) 也是 1/4

因此同一批数据仅改变 CP degree,梯度就会改变。在提供的实际 GPU 镜像中调用当前实现也复现了:

raw_metadata_T = 2
cp1_finalizer_T = 3
cp2_global_T = 2

如果整个 optimizer window 都被 mask:

  • CP=1 会静默执行一次零 loss 的 optimizer/scheduler step;
  • CP>1 的 metric denominator 为零,并且在 model.py#L1188-L1217
    (
    # Update parameters. Single optimizer.step() call handles prepare_grads, unscale,
    # clip, and inner step in one shot — avoids the double prepare_grads/unscale and
    # double grad_scaler.update that the previous external prepare_grads() flow caused.
    # In fp16 with dynamic loss scaling, step() returns (False, None, None) on overflow.
    valid_step = True
    update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
    if not getattr(args, "check_for_nan_in_loss_and_grad", True):
    # fp16 with dynamic loss scaling auto-disables this flag (see Megatron arguments.py).
    # Detect overflow via the documented (False, None, None) return signature.
    found_inf_flag = not update_successful and grad_norm is None and num_zeros_in_grad is None
    if found_inf_flag:
    valid_step = False
    current_scale = optimizer.get_loss_scale().item()
    logger.warning(
    "Inf found in gradients (step_id=%d, loss_scale=%s), skipping parameter "
    "update (dynamic loss scaling will reduce scale)",
    step_id,
    current_scale,
    )
    else:
    if isinstance(grad_norm, torch.Tensor):
    valid_step = not (torch.isnan(grad_norm) or torch.isinf(grad_norm))
    else:
    valid_step = not (math.isnan(grad_norm) or math.isinf(grad_norm))
    if valid_step:
    # Update learning rate.
    assert update_successful
    opt_param_scheduler.step(increment=args.global_batch_size)
    ) 已经执行 optimizer/scheduler 后,在
    model.py#L1237-L1259 (
    if mpu.is_pipeline_last_stage(ignore_virtual=True):
    # Average loss across microbatches.
    keys = losses_reduced[0]["keys"]
    values = None
    for x in losses_reduced:
    if values is None:
    values = x["values"]
    else:
    values += x["values"]
    assert len(keys) + 1 == values.numel()
    torch.distributed.all_reduce(values, group=mpu.get_data_parallel_group(with_context_parallel=True))
    loss_reduced = {}
    values = values.tolist()
    num_samples_or_tokens = values[0]
    for key, value in zip(keys, values[1:], strict=False):
    # No cp_size factor: num_samples_or_tokens is the all-reduced CP-local
    # token count (per-token) or sample count, so each token/sample is
    # already counted once. A `* cp_size` here would over-weight metrics by
    # CP degree under dynamic CP (and is a no-op under static CP, where the
    # count previously carried the cancelling cp factor).
    loss_reduced[key] = value / num_samples_or_tokens
    return loss_reduced, grad_norm
    ) 发生除零。

建议:

  1. CP=1 和 CP>1 使用完全相同的真实有效 token 定义,不要对每条 response 单独 clamp;Megatron finalizer 已经会对全局零 token 做安全 clamp。
  2. 对全局 T=0 的 optimizer window 在 optimizer/scheduler step 前明确 reject 或 skip。
  3. 增加以下回归测试:
    • 一条 response 全 mask;
    • 整个 window 全 mask;
    • CP1、静态 CP2 和 dynamic CP;
    • 对比真实 finalizer 之后的参数梯度或 parameter delta,而不只是 finalizer 之前的 loss.backward()。

现有 mixed-mask fixture 中每条 response 仍至少保留 14 个有效 token,因此没有覆盖这个情况:

  • test_dr_grpo.py#L548-L559 (
    def _build_loss_mask(response_length: int, variant: int) -> list[int]:
    """Tail / interior-span / strided masks, one shape per group member."""
    loss_mask = [1] * response_length
    if variant == 1:
    loss_mask[-max(response_length // 8, 1) :] = [0] * max(response_length // 8, 1)
    elif variant == 2:
    start = response_length // 3
    width = max(response_length // 8, 1)
    loss_mask[start : start + width] = [0] * width
    elif variant == 3:
    loss_mask[::8] = [0] * len(loss_mask[::8])
    return loss_mask
    )
  • test_dr_grpo.py#L847-L853 (
    def test_dr_grpo_window_fixture_matches_frozen_constants():
    window = _build_window()
    valid_token_counts = [sum(response["loss_mask"]) for response in window]
    assert len(window) == NUM_RESPONSES
    assert valid_token_counts == [1024, 896, 896, 896, 256, 224, 224, 224, 64, 56, 56, 56, 16, 14, 14, 14]
    assert sum(valid_token_counts) == FROZEN_RESPONSE_TOKENS
    )

Copilot AI review requested due to automatic review settings August 18, 2026 05:52

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.

Copilot AI review requested due to automatic review settings August 18, 2026 06:17

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.

@idvzchusvol

idvzchusvol commented Aug 18, 2026

Copy link
Copy Markdown
Author

@li126com 感谢您的 review,我们已经修复相关问题,并将相关的改动同步到PR正文中。

1. 改动

  • get_cp_local_num_tokenscp_size == 1 分支去掉逐条 clamp_min(., 1),与 cp_size > 1 统一为同一套真实有效 token 定义。
  • 对于全局 T = 0 的 optimizer window,将跳过 optimizer 和 scheduler,同时报告 warning 信息。
  • metric denominator 在 T = 0 时将报告 0,不再出现除零错误;被跳过的窗口 grad_norm 保持 0.0 而非 NaN,避免被其他检查误判。

2. 新增单测

我们新增了相关的单元测试:

  • get_cp_local_num_tokenscp_size = 1/2/4 下对全 mask response 一律计 0;
  • 整窗全 mask 时计数为 0;
  • metadata 正确标记空窗口,且该标记投递到窗口内每个 micro-batch;
  • 既有的六种 (DP, CP) 拓扑测试中加入「一条 response 全 mask」场景。

这些用例在修复前的代码上于 CP=1 全部失败、CP=2/4 通过,本次修复后全部通过。

3. 端到端验证

我们重跑了之前的端到端对照(Qwen3.5-4B / 4×H20 / TP2 / CP1 /--calculate-per-token-loss),两个 run 均完整跑完,没有NaN/Inf、OOM 或 hang。

此外,我们构造了会触发该分支的场景:训练时把部分或全部 response 的 loss_mask 置零。结果符合预期:单条 response 全 mask 时,分母精确等于真实有效 token 数;整窗全 mask 时分母为 0,跳过日志在全部 rank 触发,无除零错误。该场景需要临时改动训练代码来注入 mask,所以只作为一次性验证,测试脚本未随本 PR 提交。

Copilot AI review requested due to automatic review settings August 18, 2026 07:11

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

感谢作者持续修改。重新按题目全部验收项审查了当前 head c3f3157。此前 whole-response mask 导致 CP1/CP2 分母不一致,以及全零窗口仍执行 optimizer/scheduler 的问题已经修复;默认论文配置下的 Dr.GRPO 数学实现目前是正确的。

当前建议仍为 Request changes,只剩下面一个实现 blocker,以及一个合入前的分支同步要求。

[P1] Dr.GRPO + --normalize-advantages 在 CP 和全零 mask 下仍不正确

该选项目前在 Dr.GRPO 文档中被明确描述为支持,但实现存在三个相关问题:

  1. 静态 CP 下,advantages 和 loss masks 已经被切成 CP-local shard;whitening 却使用默认的 DP-only process group。这样不同 CP rank 会分别基于自己的 token shard 计算 mean/std,使 normalization 结果随 CP 切分变化,违反 “CP 不改变最终统计” 的验收要求。

  2. 是否进入 collective 由 rank-local 的 all_masks.numel() > 0 决定。如果某个 rank 的 local shard 为空、同组其他 rank 非空,部分 rank 会进入 all-reduce、部分 rank 会跳过,存在 distributed hang 风险。

  3. 整个 optimizer window 的 mask 都为零时,mask tensor 仍可能是非空的,因此代码会先进入 distributed_masked_whiten() 并抛出 global mask sum 为零的 ValueError。这发生在新增加的 empty-window metadata 和 optimizer/scheduler skip 之前,所以全零窗口保护在该配置下实际上不可达。

建议二选一:

  • 更简单且更符合原始 Dr.GRPO:在参数校验中拒绝 dr_grpo + --normalize-advantages,同步修改中英文文档并增加参数测试。
  • 如果继续承诺支持:静态 CP 应在 DP+CP group 上统计;dynamic CP 的 full-replicated tensor 可继续使用 DP-only group;所有 group rank 必须一致参与 collective;global valid count 为零时应统一返回 zero advantages 或跳过 whitening。

如果选择继续支持,建议至少补:

  • mixed-mask 下 CP1/CP2 whitening oracle 一致;
  • 某个 rank local shard 为空时 collective 仍能正常完成;
  • 全零 window 不抛错,并能走到 optimizer/scheduler skip。

合入前需要同步最新 main

当前 PR head 与最新 main 存在真实冲突,涉及:

  • docs/en/examples/algorithms.md
  • docs/zh/examples/algorithms.md
  • relax/backends/megatron/loss.py
  • relax/components/advantages.py
  • relax/utils/utils.py

最新 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 路径已满足:

  • 独立变体接入:存在独立的 dr_grpo 注册和 CLI estimator,与标准 GRPO 清晰区分。

  • Reward normalization:组内只减均值、不除 group std,并禁止 reward-side KL 改变原始 reward。

  • Loss aggregation:optimizer window 统计全局 response 数 N 和有效 token 数 T,使用 T/(N·B) 补偿;经过 Megatron token finalization 后得到固定分母目标 Σloss/(N·B),没有额外的 DP、CP 或 micro-batch 因子。

  • 长短 response 测试:包含 1024/256/64/16 的混合长度 oracle。

  • DP/CP、micro-batch 和 padding:已有多种 DP/CP 拓扑、micro-batch 划分、dynamic batching、THD/Bridge padding 和 fully-masked response 测试。

  • 标准 GRPO 回归:同一 oracle 同时覆盖 GRPO 与 Dr.GRPO,并验证 fixed response budget 不会改变标准 GRPO。

  • Recipe 与对比报告:相同模型、数据、batch、optimizer、sampling 和 seed,报告包含 reward、length、统一口径的 reference KL 和 grad norm/稳定性指标。

非阻塞说明

对比报告 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。除非冲突解决引入新的算法语义变化,否则没有发现其他需要阻塞合入的验收问题。

leelipeng and others added 8 commits August 18, 2026 18:08
# 🐛 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>
Copilot AI review requested due to automatic review settings August 18, 2026 15: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.

Copilot AI review requested due to automatic review settings August 18, 2026 17:46

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.

@idvzchusvol

Copy link
Copy Markdown
Author

@li126com 感谢您的 review。两个 blocker 已处理,测试增强项已补上,相关改动同步至PR正文。

1. Dr.GRPO + --normalize-advantages

根据您的建议,我们在 _validate_dr_grpo_args 中拒绝该组合,启动阶段直接报错。中英文档已同步更新,同时也新增了对应的测试用例。

2. 同步最新 main

已 rebase 到最新分支。通过 Dr.GRPO 与 RLOO 的相关测试,通过pre-commit run --all-files

3. 测试增强

我们新增以下测试用例:

  • Megatron finalizer:调用 finalize_model_gradsT = 4930 时它按 1/T 缩放梯度,结果落在 sum(loss) / (N * B)T = 0 时它不执行这次缩放(if num_tokens > 0 判断在前),梯度保持原值且有限。
  • optimizer/scheduler spy:调用 train_one_step,空窗口标记经真实 forward_step 流到跳过判断。全 mask 窗口下两个 .step() 均为 0 次、grad_norm0.0 而非 NaN;正常窗口下各 1 次。

这些用例期望测试:无有效 token 的窗口误更新参数;同样的窗口误推进 LR schedule;被跳过的窗口把 grad_norm 报成 NaN,被下游健康检查误判为发散;以及 loss 侧的窗口 scale 与 Megatron 的 1/T 归一化不再互补,使最终分母偏离 N * B

4. CP2 / full-mask smoke test

我们进行了 CP2 与 full-mask 两类 smoke test,配置对齐 docs/public/dr-grpo/experiment_manifest.json(Qwen3.5-4B / TP2 / 4×H20 / rollout 4×4 / GBS 16 / budget 4096 / lr 1e-6 / seed 1234),差异只有 context_parallel_size 1 → 2。

full-mask 沿用上一轮的做法,即临时改动训练代码注入 mask。由于这种注入方式修改了训练逻辑,同时也依赖GPU无法被 CI 复跑,因此我们暂时没有提交至本 PR 中。为了让这次验证可复查,下面给出完整的注入改动、启动方式和观测到的日志证据。

注入方式

loss_masks 生成处加一段环境变量门控的临时改动。不设该变量时行为与主线完全一致:

--- a/relax/utils/utils.py
+++ b/relax/utils/utils.py
@@ -127,6 +127,13 @@ def convert_samples_to_train_data(args, samples):
         if sample.remove_sample:
             sample.loss_mask = [0] * sample.response_length
+        # TEMPORARY DEBUG INJECTION -- not for commit.
+        # Zeroes loss masks so a real training run reaches the Dr.GRPO
+        # empty-window path (T == 0). "all" masks every response, "even" masks
+        # every other one.
+        _zero = os.environ.get("RELAX_DEBUG_ZERO_LOSS_MASK", "")
+        if _zero == "all" or (_zero == "even" and sample.index % 2 == 0):
+            sample.loss_mask = [0] * sample.response_length
         loss_masks.append(sample.loss_mask)

Dr.GRPO 的 T 就是对 loss_masks 求和,在这里置零可直接命中空窗口分支。all 置零全部 response,even 只置零一半作对照。

结果

配置 注入 Steps 结果
CP2 10/10 job succeeded;train/losstrain/grad_normtrain/ppo_klrollout/raw_rewardrollout/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 而非 NaNZeroDivisionErrorTraceback 计数为 0;partial-mask 中窗口内只有一半 response 被置零时仍有有效 token,跳过分支不应触发。实测跳过警告为 0 次,optimizer 正常执行。 三组均未出现 NaN/Inf、OOM、hang 或 Ray job failure。

Copilot AI review requested due to automatic review settings August 20, 2026 10:59

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.

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.

3 participants