Skip to content

【Task.27】refactor(algorithms): declare each algorithm once in a registry - #276

Open
Men1scus wants to merge 6 commits into
redai-studio:mainfrom
Men1scus:pr1/algorithm-registry
Open

【Task.27】refactor(algorithms): declare each algorithm once in a registry#276
Men1scus wants to merge 6 commits into
redai-studio:mainfrom
Men1scus:pr1/algorithm-registry

Conversation

@Men1scus

@Men1scus Men1scus commented Aug 14, 2026

Copy link
Copy Markdown

【Task.27】解耦算法配置:AlgorithmSpec 注册表 + 现有算法迁移

RFC: #218 · 本 PR 为两个 PR 中的第一个,不含 GDPO,可独立验证「不改数值行为」。
Base: main@4899b8f(已 rebase 到最新 main,含 RLOO #205


为什么

一个 --advantage-estimator 取值被多处独立解释:角色查表、reward 归一化、advantage 公式、policy loss 公式、REINFORCE++ 的归一化与 loss reduction,以及两轮参数校验。新增算法要把它们全部找到,漏一处的失败是晚失败——不是启动时报「算法未注册」,而是训练跑起来之后数值悄悄不对。

(早期版本的这段话举了「reinforce_plus_plus 被 argparse 接受却不在 ALGOS 里」当例子。那在本 PR 的 base 98a1274不成立:上游已经手工补进去了。这个例子来自更早的 main,写进 commit message 是我的疏忽,会诊时被指出。它恰恰说明手工维护的表会漂移,但不该被当成本 PR 基线上的现存 bug 引用。)

AlgorithmSpec 让这些事实只说一次。

关键设计

字段存字符串标识符,不存 callable。 advantage 公式在 Advantages Ray Serve deployment 里跑,policy loss 在 Megatron worker 里跑,两个进程 import 的模块子集不同。只有算法名跨进程,各自用自己的表解析。这同时让 relax/algorithms/ 不含重量级 import,注册表才能在纯 CPU runner 上测。

显式 dict 字面量,不用装饰器注册。 装饰器依赖「这个模块被 import 了吗」,而上面那两个进程的 import 图不同——一侧漏 import 就静默丢算法,正是要消灭的失败模式。

只收录有消费者的字段。 本 PR 的 8 个字段每个都有调用点:三个标识符 + kl_level/needs_full_log_probs(GSPO) + needs_critic(PPO) + requires_normalize_advantages(REINFORCE++ 两个)。GDPO 需要的另外 6 个随 PR2 一起进来,不提前占位。

本 PR 收敛到哪为止(说清楚,避免高估): reward 归一化、advantage 公式、policy loss、ALGOS 角色表、以及参数校验的算法相关部分。以及 loss.py 里按 reinforce_plus_plus{,_baseline} 选归一化方式与 loss reducer 的那两处——它们需要一个新字段(advantage_normalization),按评审意见已从 PR2 移入本 PR。没有收敛的: actor.py / critic.py / data_fields.py / controller.py / process_role 里若干 == "ppo"(都等价于 needs_critic,今天不出 bug,但第二个 value-based 算法进来时会踩)。这些不在本 PR 范围内,列在这里是为了不让读者以为「注册表已经是唯一事实源」。

角色拓扑由 needs_critic 驱动ALGOS 从注册表派生,PPO 保住 Critic。process_role 一行未改——它决定的是角色迭代顺序,属于 Controller 的编排面。

两条 advantage 路径的逐行差异

统一 handler 只覆盖 advantage 分派这一段。三处实质差异都保留在各自调用方:

差异 仅存在于 统一后如何保留
padded_total_lengths loss.py:558(maybe_padded_total_lengthsargs.qkv_format + VL/unsplit-forward 标志) 成为 handler 的 keyword-only 形参。Megatron 侧传值,Advantages deployment 无法计算、不传(保持它原有行为)。这不是装饰性的:bshd / VL / unsplit-forward 下 GAE 要按 padded 偏移切 CP 分片,传 None 不报错,只是读错 token 位置
normalize_advantages + CP mask 重建(约 80 行) loss.py:620 不进 handler,原地留在 loss.py。它在 advantage 算完之后跑,且依赖 max_seq_lens 与 CP 组——deployment 两者都没有
早退条件 loss.py:561 not mpu.is_pipeline_last_stage()components/advantages.py:152 log_probs is None and values is None and rollout_log_probs is None 不进 handler,各自留在调用方。两者语义不同(流水级 vs 数据缺失),合并需要 handler 知道自己在哪个进程里,那正是要避免的耦合

顺带回答 review 里的一个问题:max_seq_lens 没有同类问题。get_advantages_and_returns_batch 的签名里根本没有这个参数,它只服务于上表第二行的 normalize+CP 段。

等价性论证

原 RFC 依赖「ppo_utils.py 与 main 逐字节相同」。这个前提已经失效039ce87 → 98a1274 之间该文件改了 +474 −29。改用两条证据:

  1. 逐 kernel 签名核对(7 个 kernel 全查)。发现并修正一处真实漂移:get_reinforce_plus_plus_baseline_advantages 上游已移除 kl_coef 形参(KL 改由独立 k2 loss 承担),旧 adapter 仍在传。其余 6 个逐参数一致,包括 compute_policy_loss 新增的 eps_clip_c——它默认 None,main 的 loss.py 也不传,adapter 行为与之相同。
  2. characterization tests:归一化器逐常量复现原算术(1e-6 组 epsilon、--disable-grpo-std-normalization 门控),tests/algorithms/test_dispatch_parity_vs_main.py 对每个 estimator 与 ppo_utils 的裸 kernel 做逐位比对(torch.equal,不用 allclose——它的默认容差足以吞掉有偏/无偏标准差的差别)。

Ask First 声明

CLAUDE.md,以下改动落在需要事先确认的范围内,在此明确列出请维护者裁定:

  1. relax/utils/arguments.py 参数解析--advantage-estimatorchoices 由硬编码列表改为 list_algorithm_names()。取值集合本身不变(本 PR 不增不减算法)。
  2. relax/core/registry.pyALGOS:由手写 dict 改为从注册表派生。process_role 未改。这是 Controller 直接消费的编排面。
  3. 新增校验函数 validate_algorithm_args,插在 _validate_reinforce_plus_plus_args 之后——那个函数标着 frozen Task 29 contract 且其测试匹配特定文案,先跑会抢走它的错误消息。

如果其中任何一条希望换个做法,我照改。

一处刻意不做

_validate_reinforce_plus_plus_argsreinforce_plus_plus_baseline 手写了四条校验(禁 --custom-reward-post-process-path、禁 --agentic-custom-advantage-path、要求 rewards_normalization、要求组大小 > 1)。这四条逐条对应 capability flag,是注册表想消灭的重复的活样本。

本 PR 没有收编它:该函数标注为 frozen Task 29 contract,且 tests/utils/test_arguments_reinforce_plus_plus.py 直接匹配它的文案。

PR2 里收编了一半:那四条约束现在也写进了 spec 字段。原因是发现「不声明」并不中立——未声明的字段取默认值,等于 spec 主动断言了四件与实际相反的事。由于两条校验路径都是 frozen 函数先跑(主路径本来如此,YAML 路径在 PR2 里补齐),它的措辞仍然优先,测试不受影响。

RLOO(#205)的注册化迁移

本 PR 开着的期间,最新 main 合入了 RLOO。两边都不能直接取:保留本分支会丢掉 RLOO,保留 main
会把算法名分支恢复回来。所以是迁移进来,不是并排放着。

main 上的形态 迁移到
post_process_rewards 里的 == "rloo" 内联分支 group_leave_one_out normalizer
elif advantage_estimator == "rloo" 的 policy loss POLICY_LOSS_FNS["rloo"]
6 处字符串白名单里的 "rloo" reward_normalizer / advantage_fn 字段
_compute_rloo_group_diagnostics!= "rloo" 门控 reward_normalizer 判定(指标 key 前缀保持 rloo/ 不变,那是已发布的名字)
11 条启动约束 6 个 capability 字段

11 条约束压到 6 个字段,是因为其中 5 条有同一个原因:无裁剪目标没有重要性比值修正,所以
每个 optimizer step 必须正好消费产生它的那次 rollout。requires_on_policy_updates 一次性拒绝
--fully-async / --hybrid--max-staleness != 0--num-steps-per-rollout != 1
rollout_batch_size × n_samples != global_batch_size--partial-rollout /
--use-dynamic-global-batch-size。spec 里写明了 RLOO 目前是唯一成员,所以这个打包是对「下一个
无裁剪估计器」的猜测;真出现只需要其中 4 条的,应该拆字段而不是加例外。

另外 4 个字段:requires_rewards_normalizationforbids_normalize_advantages
min_group_sizeforbids_reward_side_klrequires_global_token_loss。其中
forbids_reward_side_klrequires_rewards_normalizationmin_group_size 各有两个成员
reinforce_plus_plus_baseline 也满足),不是为 RLOO 一个算法造的。

校验函数拆成 4 个而不是 1 个,原因与算法无关,是参数校验本身有推导顺序:--kl-coef 必须在
「检查 --ref-load 是否存在」之前判掉(否则真正的问题会被报成缺少参考 checkpoint),
one-update 等式必须在 global_batch_size 定稿之后判。每个函数的 docstring 写了它为什么不能挪。

拆分带来一个洞,一次会诊查出来的,已修:apply_custom_config_overrides 原先只重跑了四个里的两个,
于是一个 --custom-config-path 的 YAML 可以先切到 rloo、再设 --kl-coef--num-steps-per-rollout 4
或破坏 one-update 等式的 global_batch_size,三条约束都不会重跑。四个现在全部重跑,并抽出 derive_global_batch_size,因为
validate_batch_shape 读它写的值。变异验证:撤掉重新接上的三个调用,4 条测试变红。

第一版这个修法本身是错的,下一轮会诊查出来的。 我当时写「它们读的每个值都已定稿,所以推导顺序
在这里不适用」——对 validate_batch_shape 不成立:它读的 global_batch_size 恰恰是合并点之前
num_steps_per_rollout 派生出来的。于是一个从 grpo@4-steps 切到 rloo@1-step 的 YAML(合法配置,
从头解析应得 rollout × n = 128)被拿去和残留的 32 比较,然后拒绝。派生现已抽成
derive_global_batch_size,两条路径共用,第二次跑时跳过一致性 assert(要比较的那个值本身就是过期的)。

范围也说清楚:这个函数只堵算法校验器这一类洞。YAML 合并点之前还有六处非算法校验不重跑——
--ref-load 存在性、kl_coef/kl_loss_coef 互斥 assert、_normalize_sync_ppo_kl_args
fully-async 的 resource 检查、rollout_batch_size 派生、over-sampling assert。它们都早于本 PR 存在;
根治是把 YAML 合并挪到校验之前,比本 PR 大得多。注释里现在把这六处逐条列出,不再声称「没有洞」。

RLOO 的报错措辞里被 main 测试匹配的子串全部保留;main 的 6 个 RLOO 测试文件与 origin/main 逐字节相同,
本机 49 passed + 2 skipped(skip 只因本机无 megatron.core)。另加相对最新 main 的数值等价测试:reward normalizer 与 main 内联分支的转录逐位
比对,policy loss adapter 与 compute_rloo_loss 逐位比对,外加一条确认 adapter 忽略 ppo_kl
(无裁剪目标不该有比值项)。变异验证:把 reward_normalizer 改成 group_mean_std
policy_loss_fn 改成 ppo_cliprequires_on_policy_updates 改成 False,分别有 4 / 4 / 10
条测试变红。

loss.py 里两份重复的 REINFORCE++ 算法名(review 意见 2)

loss.py:653(advantage normalization)和 :813(loss reducer)各维护一份
{"reinforce_plus_plus", "reinforce_plus_plus_baseline"}。两者必须同步——token-global 归一化
只有配上 mask-safe reducer 才正确——但没有任何东西强制这一点。现在都读
advantage_normalization 字段。

原来的测试没抓到它,因为测试本身是瞎的:它禁的是 args.advantage_estimator in [,而实现写的
in {。现在换成覆盖全部拼法的正则(== / != / in [ / in { / in (),并给这个守卫
本身配了测试。变异验证:把 in { 放回去,测试立刻变红。

测试

1580 passed, 323 skipped        (其中 tests/algorithms 588)

失败集合与 main@4899b8f 在同一台机器上逐条一致(2 failed + 2 errors:/dev/shm 在 macOS
不存在,以及一个既有的 test_reward_router 失败)。验证方式:在该 commit 上开一个 detached
worktree 重跑同样的用例,得到同一组失败。(早期版本说是用 git stash 验证的——改动已经提交,
stash 什么也不会动,那个说法是错的。)

pre-commit run --all-files 全部通过。

运行时证据(Modal,H100)

CPU 测试证明的是映射表和数值;Controller 真的据此把服务建起来、算法真的按注册表跑这一段
只能靠跑。四次冒烟,全部 succeeded

配置 revision 证据
1×H100 GDPO d11e247 raysubmit_Lap4dJNZgGvFqyjw,8 步,loss 0.403/0.091/0.138/0.219
2×H100 GDPO d11e247 raysubmit_TGe58MhMv62FZTb7,4 步,SeqlenBalancedSampler with dp_size=2
1×H100 PPO d11e247 raysubmit_9p3bUqYts8kkZ6ci,critic 服务 + 8 步 value_loss
1×H100 RLOO a61b68f raysubmit_HR4gXykqDnKtXC7X,4 步,pg_clipfrac 恒为 0

revision 归属,如实写出来而不是含糊成「都在 HEAD」:

  • 这四个 revision 都在 PR2 的堆叠分支上,不是本 PR 的 head。两次 GDPO 冒烟验证的是
    PR2 的代码;对本 PR 有直接意义的是 PPO 与 RLOO 两次。
  • 后续 rebase 换掉了这些 SHA。相对冒烟时的 a61b68f,当前 PR2 tip 在 relax/ 下只有一处
    可执行改动:observed_reward_signalexcept 元组从 (TypeError, ValueError) 加宽为
    也含 KeyError, IndexError(仅指标路径)。arguments.py 的非注释改动为零——YAML 校验补齐
    本来就在冒烟的那一版里,rebase 只是把它从 PR2 移到了本 PR。

PPO 是本 PR 改动最大的既有算法——唯一 needs_critic=True 的那个,ALGOS 里的 critic 条目
从写死的名字比较改成由该字段推导:

Deploying new version of Deployment(name='Critic', app='critic')
Adding 1 replica to Deployment(name='Critic', app='critic')
(ServeReplica:critic:Critic pid=3603)                    ← replica 真的起来了
critic value_loss  6.77, 8.39, 4.22, 8.76, 7.42, 3.97, 2.93, 4.63

RLOO 是本 PR 重写最多的那个(reward normalizer、policy loss adapter、11 条启动约束压成
6 个 capability 字段)。RLOO 本身在上游 #20598a7234)已有端到端记录——那条 commit message 写了 2-GPU canary 与 GRPO/RLOO 各 60/60 步的配对实验;这里跑的是注册化之后的这一版。三条独立证据说明跑的确实是 RLOO 而不是
退回 PPO-clip:

train/pg_clipfrac   0.0 每一步   ← main 自己的 docstring 称它为
                                   "a wiring self-check that the unclipped path is active"
train/loss          -0.032 / -0.048 / -0.028   ← REINFORCE 目标的符号,不是裁剪代理
rollout/rloo/baseline_mean   0.344 → 0.656
rollout/rloo/adv_abs_mean    0.259 → 0.518     ← 把 `!= "rloo"` 名字门控换成按
                                                  `reward_normalizer` 判之后,诊断照常发布
train/grad_norm     1.14 → 0.73 → 0.44

启动校验一次误拦都没有。冒烟脚本从已知能跑的 GDPO 那份派生,只换算法相关的一段——镜像、数据准备、
并行度全部不变,变的正好是被测的东西。

2 卡那次dp_size=2)是唯一能验到跨 rank 归约的:单卡下 DP world size 是 1,所有 all-reduce
都是空操作。

另外 advantage_gae 此前只有 co_names 检查——它看不出丢参数、KL 系数符号翻转、或者终止
奖励注入到错误的 token。现在与 main 内联分支的转录做逐位数值比对
test_gae_adapter_reproduces_mains_reward_shaping),并配一条反向测试确认 adapter 真的
读了 kl_coef / gamma / lambd,免得转录和实现一起错还互相同意。

仍然没有运行证据的

如实列出,不当成已验证:fully-async 模式(GDPO 被参数校验挡在外面,但既有算法在该模式下
走的是同一套新分派)、多机PP > 1。这些靠读代码关不掉,只能靠跑。

@li126com

Copy link
Copy Markdown
Member

整体方向是正确的:AlgorithmSpec、adapter table 和统一 dispatch 已经显著减少了算法接入时需要修改的位置,现有算法 kernel 也基本保持不变。focused CPU tests 共 554 passed,当前 head 的lint/pre-commit checks 也是绿色。

不过当前版本还不能作为独立的注册与分发基础 PR 合入,主要有以下两项阻塞:

  1. 需要基于最新 main 重新集成 RLOO。

    最新 main 已新增 RLOO,包括 reward-side LOO (https://github.com/redai-infra/Relax/blob/98a72349c7d0368440eb6b0c6849e9d0f2ba8cef/relax/utils/utils.py#L185-L223)、advantage 与 unclipped policy loss (https://github.com/redai-infra/Relax/blob/98a72349c7d0368440eb6b0c6849e9d0f2ba8cef/relax/backends/megatron/loss.py#L579-L583) 以及对应 CLI 配置,但本 PR 的 registry/spec 和 parity tests 都不包含 RLOO。目前不能简单选择冲突的一侧:保留 PR 会丢失 RLOO,保留 main 又会恢复算法名分支。

    请先 rebase 到最新 main,将 RLOO 的 reward normalization、advantage、unclipped policy objective 和参数约束迁入注册机制,并补充相对最新 main 的数值等价测试。

  2. 基础迁移仍存在重复维护的算法名称集合。

    loss.py 仍在 advantage normalization (https://github.com/redai-infra/Relax/blob/801d881230e421d562736b2afa04d76ad3273d82/relax/backends/megatron/loss.py#L653-L656) 和 loss reducer (https://github.com/redai-infra/Relax/blob/801d881230e421d562736b2afa04d76ad3273d82/relax/backends/megatron/loss.py#L813-L816) 两处重复维护 REINFORCE++ 算法名。现有测试只禁止 in [,因此没有发现实现使用的是 in {。【Task.27】feat(algorithms): add GDPO multi-reward advantage estimator #277 已通过 advantage_normalization capability 修复这两处,但这是现有算法的通用注册迁移,建议移到 【Task.27】refactor(algorithms): declare each algorithm once in a registry #276 中完成。

完成上述迁移、更新 parity tests,并重新运行 CPU tests 和 pre-commit run --all-files 后,这个基础 PR 的架构方向可以接受。

@Men1scus
Men1scus force-pushed the pr1/algorithm-registry branch from 801d881 to a3ebe9d Compare August 24, 2026 03:26
Copilot AI lite review requested due to automatic review settings August 24, 2026 03:26

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

Refactors algorithm metadata and dispatch into a declarative AlgorithmSpec registry while preserving existing behavior.

Changes:

  • Centralizes algorithm capabilities, validation, roles, rewards, advantages, and policy-loss dispatch.
  • Adds parity and characterization tests.
  • Updates documentation and examples.

Reviewed changes

Copilot reviewed 25 out of 26 changed files in this pull request and generated 7 comments.

Show a summary per file
File Review status
tests/algorithms/test_reward_normalizers.py Reviewed; no findings.
tests/algorithms/test_post_process_rewards_dispatch.py Reviewed; no findings.
tests/algorithms/test_policy_loss_dispatch.py Reviewed; no findings.
tests/algorithms/test_dispatch_parity_vs_main.py Reviewed; no findings.
tests/algorithms/test_arguments_spec_driven.py Reviewed; no findings.
tests/algorithms/test_algos_roles.py Reviewed; no findings.
tests/algorithms/test_algorithm_registry.py Reviewed; no findings.
tests/algorithms/test_advantage_estimators.py Reviewed; no findings.
tests/algorithms/__init__.py Reviewed; no findings.
relax/utils/utils.py Reviewed; no findings.
relax/utils/metrics/metric_utils.py Reviewed; no findings.
relax/utils/arguments.py Moderate (3 votes): Re-run reward-side-KL, update-schedule, and batch-shape validation after YAML estimator overrides, or reject such overrides.
relax/core/registry.py Reviewed; no findings.
relax/components/advantages.py Nit (2 votes): Clarify the estimator-registry lookup in the docstring.
relax/backends/megatron/loss.py Reviewed; no findings.
relax/algorithms/spec.py Reviewed; no findings.
relax/algorithms/rewards.py Reviewed; no findings.
relax/algorithms/policy.py Reviewed; no findings.
relax/algorithms/advantages.py Reviewed; no findings.
relax/algorithms/__init__.py Reviewed; no findings.
examples/generate_reward_model/post_process_genrm_swap.py Reviewed; no findings.
docs/zh/guide/configuration.md Nit (3 votes): Add rloo to the supported estimator list.
docs/zh/guide/adding-an-algorithm.md Nit (3 votes): Remove or add the nonexistent relax/algorithms/numerics.py entry.
docs/en/guide/configuration.md Nit (3 votes): Add rloo to the supported estimator list.
docs/en/guide/adding-an-algorithm.md Nits (3 votes each): Remove or add the nonexistent relax/algorithms/numerics.py entry and fix the duplicated “components” wording.
docs/.vitepress/config.mts Reviewed; no findings.
Suppressed comments (3)

relax/algorithms/spec.py:43

  • These string-valued capability fields are consumed through equality checks, but they are not validated at startup. A typo such as kl_level="sequnce" or advantage_normalization="token_globa" would silently select the token-KL or ordinary-whitening path; _assert_spec_implementations_resolve only validates the three table identifiers. Validate the allowed values in AlgorithmSpec or the startup validator so a malformed registry entry fails fast.
    kl_level: str = "token"
    """``"token"`` or ``"sequence"``; GSPO constrains the sequence as a whole."""

    advantage_normalization: str = "whiten"

relax/utils/arguments.py:3147

  • loss_type is accepted from YAML, but is_sft and the SFT/RL setup were computed before this helper runs. If an RL invocation overrides it to sft, this early return leaves the outer validation in RL mode while the controller later selects SFT roles; the reverse override leaves SFT-only flags in an RL run. Reject loss_type in custom overrides or rerun the complete mode-dependent validation with a recomputed mode.
    if args.loss_type in ("sft", "sft_loss", "sft-loss"):
        return

relax/utils/utils.py:432

  • is_group_normalized is implemented as reward_normalizer != "none", so this gate will treat every future non-none normalizer as group-based. A batch-level or per-sample normalizer would therefore trigger group-preserving debug subsampling and receive only selected groups, even though it does not require complete groups. Add an explicit normalization-scope capability (or otherwise restrict this check to group normalizers).
            args.custom_reward_post_process_path is None
            and get_algorithm(args.advantage_estimator).is_group_normalized
            and args.rewards_normalization

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread docs/en/guide/adding-an-algorithm.md Outdated
├── rewards.py reward normalization strategies + REWARD_NORMALIZERS
├── advantages.py advantage estimators + ADVANTAGE_FNS
├── policy.py policy loss adapters + POLICY_LOSS_FNS
└── numerics.py shared numeric constants and degeneracy guards

The output must be **one scalar per sample**. That constraint is what keeps the
TransferQueue schema fixed — an algorithm reading several reward components collapses them
components to a scalar here.
Comment thread docs/en/guide/configuration.md Outdated
| Parameter | Type | Default | Options | Description |
|-----------|------|---------|---------|-------------|
| `--advantage-estimator` | str | grpo | `grpo`, `gspo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`, `ppo`, `sapo`, `cispo` | Advantage estimator. OPD is independent of this choice; enable it with `--use-opd` and its KL/loss coefficient |
| `--advantage-estimator` | str | grpo | generated from `ALGORITHM_SPECS` in `relax/algorithms/spec.py`; currently `grpo`, `gspo`, `sapo`, `cispo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline` | Advantage estimator. `--help` is authoritative: the choices are read from the registry, so a new algorithm appears there without this table being edited. OPD is independent of this choice; enable it with `--use-opd` and its KL/loss coefficient |
Comment thread docs/zh/guide/adding-an-algorithm.md Outdated
├── rewards.py reward 归一化策略 + REWARD_NORMALIZERS
├── advantages.py advantage 估计器 + ADVANTAGE_FNS
├── policy.py policy loss 适配器 + POLICY_LOSS_FNS
└── numerics.py 共享的数值常量与退化判定
Comment thread docs/zh/guide/configuration.md Outdated
| 参数 | 类型 | 默认值 | 可选值 | 说明 |
|------|------|--------|--------|------|
| `--advantage-estimator` | str | grpo | `grpo`, `gspo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`, `ppo`, `sapo`, `cispo` | 优势估计器。OPD 独立于该选项,使用 `--use-opd` 及对应 KL/loss 系数启用 |
| `--advantage-estimator` | str | grpo | 由 `relax/algorithms/spec.py` 的 `ALGORITHM_SPECS` 生成,当前为 `grpo`、`gspo`、`sapo``cispo`、`ppo`、`reinforce_plus_plus`、`reinforce_plus_plus_baseline` | 优势估计器。以 `--help` 为准:取值直接读注册表,新增算法无需改这张表即可出现。OPD 独立于该选项,使用 `--use-opd` 及对应 KL/loss 系数启用 |
Comment on lines +121 to +123
`rollout_data`, computes KL divergences, then delegates to the estimator
the algorithm registry names for `self.config.advantage_estimator`
(see `relax.algorithms.advantages.ADVANTAGE_FNS`).
Comment thread relax/utils/arguments.py
Comment on lines +3154 to +3155
_validate_reinforce_plus_plus_args(args, is_sft=False)
validate_algorithm_args(args)
`--advantage-estimator` was interpreted independently in six places: role
lookup, reward normalisation, the advantage formula, the policy loss, and two
rounds of argument validation. Adding an algorithm meant finding all of them,
and missing one failed late -- `reinforce_plus_plus` was accepted by argparse
while absent from `ALGOS`, crashing in `controller.register_all_serve`.

`AlgorithmSpec` states those facts once. It holds string identifiers rather
than callables: the advantage formula runs in the `Advantages` Ray Serve
deployment while the policy loss runs in the Megatron worker, and those two
processes import different module subsets, so each resolves the name against
its own table. That also keeps the module free of heavy imports, which is what
lets the registry be tested on a CPU-only runner.

Only fields with a consumer are included. Role topology comes from
`needs_critic`, so PPO keeps its Critic and nothing hard-codes the name;
`process_role` is untouched, since it selects the role *iteration order* and
that is the controller's orchestration surface.

Both advantage call sites now share one handler while keeping their real
differences: the Megatron path passes `padded_total_lengths` (GAE slices CP
shards at padded offsets, and passing nothing there reads the wrong token
positions rather than raising), the Advantages deployment cannot compute it and
passes nothing.

RLOO (redai-studio#205), which landed on main after this branch opened, is migrated in
rather than merged alongside: keeping either side of that conflict was wrong,
since taking this branch drops RLOO and taking main restores the algorithm-name
branches. Its reward stage becomes the `group_leave_one_out` normaliser, its
unclipped objective a `POLICY_LOSS_FNS` entry, and its eleven startup
constraints six capability fields. `requires_on_policy_updates` is one field
for five of those knobs because they have one cause -- an objective with no
importance-ratio correction cannot account for the policy having moved -- and
the spec says so, including that RLOO is currently its only member.

Two duplicate REINFORCE++ name sets in `loss.py` (advantage normalisation at
main's 691, the loss reducer at 851) become `advantage_normalization`. They had to
stay in step because token-global normalisation is only correct together with
the mask-safe reducer, and nothing enforced that. The test that was supposed to
catch them was blind: it banned `args.advantage_estimator in [` while the
implementation wrote `in {`. It is now a regex over every spelling, with its
own test, and reintroducing `in {` turns it red.

No behaviour change is intended: the normalisers reproduce the previous
arithmetic constant for constant, including the 1e-6 group epsilon and the
`--disable-grpo-std-normalization` gate, and RLOO's reward output is compared
against a transcription of main's inline branch rather than against the helper
it shares.

`apply_custom_config_overrides` re-runs every algorithm validator, not two of
them. `validate_reward_side_kl`, `validate_update_schedule` and
`validate_batch_shape` were split out of `validate_algorithm_args` because
validation has a derivation order, and only two of the four were wired back
into the override path -- so a YAML file could select rloo and then set
`--kl-coef`, `--num-steps-per-rollout 4`, or a `global_batch_size` that breaks
the one-update guarantee, with nothing objecting. `derive_global_batch_size` is
extracted for the same reason: `validate_batch_shape` reads the value that
derivation writes, and re-running the validator without it rejected a
legitimate config (a YAML moving 4 steps to 1 got compared against the batch
size derived from 4). The comment above the calls lists what this still does
*not* cover -- six non-algorithm checks that run before the merge -- because
closing that class means merging the YAML before validation, which is larger
than this change.

Provenance in `test_dispatch_parity_vs_main.py` was wrong and is corrected:
the header named a main SHA that does not exist in the repository, `MAIN_SHA`
held a third, unrelated commit, and the transcribed line numbers pointed at a
July revision with no rloo in it. All of them now name
main@4899b8f3a90489840a736897b4c341d87c6267cf and its actual lines; nothing
between 98a7234 and that commit touched advantages.py, loss.py, utils.py or
ppo_utils.py.

Docs: the estimator table listed every registered algorithm except the one
this branch adds, and the module tree named `numerics.py`, which does not
exist here.

Tests: 1580 passed, 323 skipped. The 2 failures + 2 errors are identical to
main@4899b8f on this machine (no /dev/shm on macOS; one pre-existing
reward_router failure), verified in a detached worktree at that commit.
`pre-commit run --all-files` clean.
Copilot AI review requested due to automatic review settings August 24, 2026 04:32
@Men1scus
Men1scus force-pushed the pr1/algorithm-registry branch from a3ebe9d to 309cd81 Compare August 24, 2026 04:32
@Men1scus

Copy link
Copy Markdown
Author

两条阻塞都改完了,已 rebase 到 main@4899b8f,当前 head 是 309cd81。因为是 force-push,您之前评论里指向的代码位置可能对不上了。

1. RLOO 挪进注册表了

  • reward 归一化:原来在 post_process_rewards 里写死的 == "rloo" 分支,现在是一个叫 group_leave_one_out 的 normalizer
  • policy loss:原来的 elif == "rloo",现在是 POLICY_LOSS_FNS 里的一项
  • 指标:_compute_rloo_group_diagnostics 原来靠算法名 != "rloo" 决定要不要算,现在看这个算法用的是不是 group_leave_one_out 归一化。rloo/ 这个指标前缀没动,它已经发布出去了
  • 启动校验:原来散在 arguments.py 里的 11 条 == "rloo" 检查,现在是 spec 上的 6 个字段

这 11 条里有 5 条(禁 fully-async/hybrid、--max-staleness 0--num-steps-per-rollout 1、一次 rollout 只做一次更新、禁 partial-rollout 和动态 batch)合成了一个字段 requires_on_policy_updates,因为它们是同一个原因:RLOO 的目标函数没有重要性采样比值这一项,所以策略一动就没法修正。spec 里注明了目前只有 RLOO 用这个字段。

main 带过来的 6 个 RLOO 测试文件我一个字没改(git diff origin/main 是空的),本机跑 49 通过 + 2 跳过,跳过的两条只是因为我这台机器没装 megatron.core。另外加了一组和 main 对比的测试:把 main 里那段内联代码抄下来,喂同样的输入,要求结果完全相同(用 torch.equal,不是 allclose)。反过来验证过——把 spec 里的 reward_normalizer 改成 group_mean_stdpolicy_loss_fn 改成 ppo_cliprequires_on_policy_updates 改成 False,分别有 4 / 4 / 10 条测试变红。

2. loss.py 里那两处重复的名字列表

:659:819,现在都读 advantage_normalization 这个字段。

您说得对,原来那条守卫测试是瞎的:它禁的写法是 advantage_estimator in [,而实现里写的是 in {,所以从来没拦住过。现在换成一条正则,==!=in [in {in ( 都能抓到,并且给这条正则本身也写了测试。把 in { 放回代码里,它立刻变红。

另外修了 Copilot 在新 head 上提的三条

这三条单看本 PR 也确实成立:

  • apply_custom_config_overrides 只重跑了 4 个校验函数里的 2 个。也就是说,一个 --custom-config-path 的 YAML 可以先把算法切成 rloo,再设 --kl-coef--num-steps-per-rollout 4、或者一个破坏「一次 rollout 一次更新」的 global_batch_size,而没有任何检查会拦。现在 4 个全跑,并且把 global_batch_size 的推导单独抽成了 derive_global_batch_size——因为 validate_batch_shape 读的正是这个推导写出来的值,只重跑校验不重跑推导会误拦合法配置。撤掉这三个新加的调用,4 条测试变红
  • 文档的模块树里写了 numerics.py,那个文件是 【Task.27】feat(algorithms): add GDPO multi-reward advantage estimator #277 才有的
  • 配置文档的算法列表漏了本 PR 新增的 rloo

顺带更正 test_dispatch_parity_vs_main.py 的出处标注:文件头引用的 main SHA 在仓库里根本不存在(是我把短 SHA 补成 40 位时编的),MAIN_SHA 常量又是另一个不相干的 commit,而抄下来的那些行号出自一个七月的 revision——那时候 rloo 还没进 main。现在统一标成 main@4899b8f3a904… 和它的真实行号。

验证

全量 1580 passed, 323 skipped。失败的 2 个 + 报错的 2 个,和 main@4899b8f 在我这台机器上是同一组(macOS 没有 /dev/shm,加一个既有的 test_reward_router 失败),我在那个 commit 上单独开了个 worktree 跑过对照。pre-commit run --all-files 通过,GitHub CI 五项全绿。

顺序按您说的来:先审这条。合了之后我把 #277 rebase 上去,那边的 diff 就只剩 GDPO 了。

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 25 out of 26 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

docs/en/guide/adding-an-algorithm.md:114

  • The sentence has an extra "components" and is grammatically incorrect; it should say that an algorithm reading several reward components "collapses them to a scalar here."
TransferQueue schema fixed — an algorithm reading several reward components collapses them
components to a scalar here.

relax/algorithms/advantages.py:117

  • This extracted PPO adapter has no execution-level parity coverage: the registry test only checks that get_advantages_and_returns_batch appears in co_names, while the existing GAE tests call ppo_utils directly. A dropped/reordered kl_coef, misplaced terminal reward, or omitted padded_total_lengths would therefore pass despite changing PPO (especially CP/VL) behavior. Add a test through compute_advantages_and_returns that compares both outputs with and without a non-None padded layout.
    return get_advantages_and_returns_batch(
        total_lengths,
        response_lengths,
        values,
        shaped_rewards,

Comment thread relax/algorithms/spec.py
Comment on lines +40 to +43
kl_level: str = "token"
"""``"token"`` or ``"sequence"``; GSPO constrains the sequence as a whole."""

advantage_normalization: str = "whiten"
@li126com

Copy link
Copy Markdown
Member

这轮相较上次进展很大:RLOO 已迁入 reward、advantage、policy loss 和参数能力注册;loss.py 两处重
复的 REINFORCE++ 名称集合也已经由 advantage_normalization 驱动。当前 tests/algorithms 为 588
passed,GitHub Python 3.10/3.11/3.12、Lint 和 Pre-commit 均通过。

不过,严格按照 Task 27 的验收要求,当前仍建议 Request changes。

1. [验收阻塞] critic/value 能力尚未真正由注册表驱动

needs_critic 当前主要用于生成 ALGOS,但实际训练编排和 value 数据流仍在多处按 "ppo" 判断:

  • core/registry.py:128

(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/core/registry.py#L128)

  • core/controller.py:107

(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/core/controller.py#L107)

  • backends/megatron/actor.py:790

(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/backends/megatron/actor.py#L790)

  • utils/training/data_fields.py:35

(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/utils/training/data_fields.py#L35)

  • components/critic.py:123

(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/components/critic.py#L123)

因此,注册第二个 needs_critic=True 的算法后,虽然 parser 和 ALGOS 会接受它,但 values 字段、
critic 等待/回传以及 actor 侧 GAE 数据处理仍不会自动接通。这正是本题希望消除的“新增算法需要继续
寻找散落分支”的风险。

建议让这些调用点读取准确的 capability 或 handler identifier;如果部分行为并不等同于
needs_critic,应增加更精确的能力字段,而不是继续比较 PPO 名称。建议补一个临时注册的第二 critic
算法测试,验证 role、data fields 和 critic/value 流程不需要新增算法名判断。

2. [P1] custom YAML 中显式设置的 global batch 会被静默覆盖

apply_custom_config_overrides()
(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/utils/arguments.py#L3197)
在合并 YAML 后无条件执行:

derive_global_batch_size(args, enforce_consistency=False)

只要 num_steps_per_rollout 非空,YAML 显式设置的 global_batch_size 就会先写入、随后被推导值静默
覆盖。这样既没有采用用户配置,也没有报告冲突,与文档“YAML key 覆盖已有参数”的契约不一致,并可能
改变训练 batch/budget。

建议区分 YAML 是否显式包含 global_batch_size:

  • YAML 只修改推导输入时,可以重新推导;
  • YAML 显式提供 global_batch_size 时,应保留并校验,冲突则 fail-fast。

请补 num_steps_per_rollout != None 且 YAML 显式设置 global_batch_size 的回归测试。

3. [P2] PPO/GAE adapter 缙少执行级数值等价测试

advantage_gae
(https://github.com/redai-infra/Relax/blob/309cd81912bbfbe6be43a5daaaac6bb3345000d2/relax/algorithms/advantages.py#L81)
承接了 PPO 的 reward shaping、terminal reward、GAE 参数和
padded_total_lengths,但当前测试主要确认 bytecode 中出现了 kernel
名称。该检查无法发现漏传参数、参数错位或 padded CP 布局丢失。

题目明确要求现有算法数值等价性。建议通过统一 dispatcher 与 main 原逻辑逐位对拍,至少覆盖:

  • padded_total_lengths=None;
  • 非空 padded layout;
  • 非零 kl_coef、terminal reward、gamma/lambda。

另外,kl_level 和 advantage_normalization 是有限枚举但运行时没有 fail-fast 校验,拼错会静默选择
默认公式路径。建议在 AlgorithmSpec.post_init 或启动校验中检查合法取值,并补非法值测试。

完成这些后,#276 才真正达到“算法名称、能力和实现通过注册表管理”的验收目标。

…not the name

`needs_critic` reached `ALGOS`, and stopped there. Everything downstream of the
role table still asked whether the estimator was literally `"ppo"`, so
registering a second value-based algorithm would have been accepted by argparse
and by `ALGOS` and then quietly not switched on any of the value plumbing: no
critic role walked, no critic placement group, `values` left on CPU, the critic
consumer handed the wrong rollout fields, the critic never waiting for data.
That is the class of bug the registry exists to remove, one layer further in
than it had reached.

Eight call sites, not the five the review listed. Three more compare with `!=`
and do not turn up in a search for `== "ppo"`:

    core/registry.py:128            role topology
    core/controller.py:107          critic co-hosted on the actor's PG
    backends/megatron/actor.py:790  critic's `values` moved to GPU
    backends/megatron/actor.py:823  who computes GAE under fully_async
    backends/megatron/actor.py:2250 `_put_critic_values_to_transfer_queue`
    components/critic.py:123        the critic's own wait loop
    utils/training/data_fields.py   the critic consumer's field set
    utils/training/ppo_utils.py:21  the `--resource` critic entry check

`algorithm_needs_critic(config)` reads the spec rather than `args.use_critic`,
which already carries the same answer: `use_critic` is only set once
`validate_algorithm_args` has run, and `process_role` and the controller's
placement logic read a config that may not have been through it. Reading the
registry makes the answer independent of call order. Unknown or absent
estimators answer False, because SFT and the debug-only role paths reach these
sites with no estimator at all.

Two `== "ppo"` checks are deliberately left: `_compute_zero_std_metrics` in
`distributed/ray/rollout.py` and `agentic/rollout.py` asks whether one prompt
has several responses, which is not the critic capability and is not
`min_group_size` either -- grpo has `min_group_size=1` and is group-based. That
needs a field of its own and is not this change.

The test registers a second `needs_critic=True` spec and asserts the role
topology, the rollout fields and the resource check all follow it. It compares
topology *identity*, not member names: every role set carries a `critic` member
and `ALGOS` is what filters it, so the first draft passed while reading the
non-critic topology. Reverting any of the three converted sites turns it red.
…r it

`apply_custom_config_overrides` merges the YAML and then re-derives
`global_batch_size` from `num_steps_per_rollout`. With `num_steps_per_rollout`
set, a YAML file naming `global_batch_size` had its value written by the merge
loop and replaced one statement later, so the run used neither the configured
number nor an error -- the single outcome the "a YAML key overrides the
argument" contract rules out. It also moves the training batch and the token
budget without saying so.

`enforce_consistency=False` was right for the case it was added for and is kept:
a YAML that switches `num_steps_per_rollout` from 4 to 1 must get
`rollout * n`, and the pre-merge value is stale by construction, so comparing
against it rejects a legitimate config. What the call could not distinguish is
where the current value came from. `data` already knows: a key the YAML names
is an intent, a key it does not name is a leftover. Derive first so the error
can quote both numbers, then refuse the conflict rather than picking a winner.

Three tests: the conflict is refused, a YAML that names the value the derivation
would reach anyway still passes, and the re-derivation the `enforce_consistency`
flag exists for still happens. Only the first is a mutation target -- the second
guards against over-triggering and the third against breaking the original fix.
… its bytecode

`advantage_gae` had only a `co_names` check, which asserts the kernel's name
appears in the adapter's bytecode. That check survives every way this adapter
can actually be wrong, and it is the adapter with the most to get wrong: it
shapes the reward in place before delegating, and it carries
`padded_total_lengths`, the one argument main's two call sites disagreed on.

Four mutations, all of which the old check passes and these fail:

    padded_total_lengths dropped        -> reads the wrong token positions
    kl_coef sign flipped                -> KL pushes the wrong way
    terminal reward dropped             -> trains on KL alone
    gamma and lambd swapped             -> wrong discounting

main's PPO branch is transcribed from components/advantages.py:181-193 and the
megatron duplicate at loss.py:585-602 rather than regenerated, per this file's
existing rule -- regenerating turns the comparison into the implementation
checking itself. Each side gets its own tensors: `advantage_gae` mutates `kl`
in place (`k *= -args.kl_coef`), so sharing them would make the second call
read already-shaped rewards.

The `cp_disabled` fixture that made the reinforce++ adapter testable does the
same here. That fixes the reachability problem and creates a coverage boundary
worth stating: `padded_total_lengths` is only *consumed* when `cp_size > 1`, and
these run at 1. So the numeric tests pin the values and the spy pins that the
argument survives the adapter in the right keyword and position -- the padded
slicing itself needs a real context-parallel group and is not covered here. The
test says so rather than leaving the gap to be inferred.
…ation

`kl_level` and `advantage_normalization` are enum-like strings consumed by
equality checks -- `advantage_normalization == "token_global"` at loss.py:659
and 819, `kl_level == "sequence"` at loss.py:919. Every other string takes the
else branch, so `"token-global"` or `"Sequence"` in a registry entry does not
fail: the run starts, trains, and uses a different formula than the one the
spec meant to select.

The other spec fields do not have this problem, which is why these two were
missed. `advantage_fn` and `policy_loss_fn` are dictionary keys, so a typo
raises a KeyError -- and `_assert_spec_implementations_resolve` already pulls
that failure forward to startup so it names the culprit instead of surfacing
inside a worker. These two needed the same treatment and had none.

`__post_init__` rather than a startup validator, because `ALGORITHM_SPECS` is a
module-level literal: the check runs at import, so a bad entry cannot reach a
worker, let alone a training step. The allow-lists sit next to the class with
the call sites they mirror named in the comment, and a test asserts the shipped
specs stay inside them, so the lists cannot drift away from the registry they
guard.
`docformatter` runs in pre-commit but not in the ruff pass I was checking
locally, so four summary lines went in too long and the hook rewrapped them --
splitting `k *= ...` and `(`k *= ...`)` across a line break in the process.
Shortening the summaries is the fix that keeps both the enforced format and a
readable first line, rather than committing the wrap.
@Men1scus

Copy link
Copy Markdown
Author

三项都已处理,head bae4ced,在 309cd81 上追加,没有 rebase,您之前评论指向的代码位置仍然有效。CI 五项全过。

1. [验收阻塞] critic 能力由注册表驱动

新增 algorithm_needs_critic(config) 读 spec,转换了 8 处——您列的 5 处之外还有三处,其中两处是 != "ppo" 形式,搜 == "ppo" 找不到:

  • backends/megatron/actor.py:823(fully_async 下由谁算 GAE)
  • backends/megatron/actor.py:2250(_put_critic_values_to_transfer_queue)
  • utils/training/ppo_utils.py:21(启动时要求 --resource 有 critic 条目)

没有复用 args.use_critic:它承载同一答案,但只有在 validate_algorithm_args 跑过之后才存在,而 process_role 和 controller 的 placement 逻辑读到的 config 可能还没经过它。读 spec 让答案不依赖调用顺序。

已补您要的测试:临时注册第二个 needs_critic=True 的 spec,断言 role 拓扑、rollout 字段和 --resource 校验都跟着走。拓扑比的是枚举身份而非成员名——每个 role set 都带 critic 成员、靠 ALGOS 过滤,按成员名比会假通过,第一版就是这么错的。撤回任意一处转换,这条测试变红。

有两处 == "ppo" 我故意保留:_compute_zero_std_metrics(distributed/ray/rollout.py:4054agentic/rollout.py:1296)问的是「一个 prompt 是否多响应」,不是 critic 能力;min_group_size 也区分不了,grpo 是 1 却确实分组。这正是您说的「应增加更精确的能力字段」,我认为该单独一个 PR,没有混进本次改动。

2. [P1] YAML global_batch_size 不再被静默覆盖

data 这个 dict 本来就知道 YAML 说了什么:显式命名 global_batch_size 是意图,没命名才是推导前留下的过期值。现在显式命名时冲突 fail-fast,错误信息同时报出 YAML 的值和推导值;只改推导输入时仍照旧重新推导,enforce_consistency=False 当初要解决的问题没有回退。

三条回归测试:冲突被拒、一致值存活、num_steps_per_rollout 4→1 仍重新推导。

3. [P2] GAE 数值等价与枚举校验

对拍:替掉 co_names 检查,从 main 转录原逻辑(components/advantages.py:181-193loss.py:585-602)逐位比对,覆盖 kl_coef 0 与 0.05、gamma/lambd 1.0 与 0.99+0.95、terminal reward。四种变异实测:丢 padded_total_lengthskl_coef 符号反、丢 terminal reward、gamma/lambd 位置对调——旧检查全部放行,新测试全部拦下

一个覆盖边界要主动说明:padded_total_lengths 只在 cp_size > 1 时被消费(ppo_utils.pyall_gather_with_cp 分支),而这些测试跑在 cp_size=1。所以我测到的是「参数原样传到 kernel、keyword 与位置都对」,不是 padded 切片本身正确——后者需要真的 context-parallel 组。测试 docstring 里写明了这一点。

枚举校验:放在 AlgorithmSpec.__post_init__ 而不是启动校验。ALGORITHM_SPECS 是模块级字面量,所以校验在 import 期执行,拼错的 spec 连 worker 都到不了。四条负向测试(大小写错、缩写、连字符、无效值),外加一条断言允许值列表不与注册表漂移。

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