【Task.27】refactor(algorithms): declare each algorithm once in a registry - #276
【Task.27】refactor(algorithms): declare each algorithm once in a registry#276Men1scus wants to merge 6 commits into
Conversation
|
整体方向是正确的:AlgorithmSpec、adapter table 和统一 dispatch 已经显著减少了算法接入时需要修改的位置,现有算法 kernel 也基本保持不变。focused CPU tests 共 554 passed,当前 head 的lint/pre-commit checks 也是绿色。 不过当前版本还不能作为独立的注册与分发基础 PR 合入,主要有以下两项阻塞:
完成上述迁移、更新 parity tests,并重新运行 CPU tests 和 pre-commit run --all-files 后,这个基础 PR 的架构方向可以接受。 |
801d881 to
a3ebe9d
Compare
There was a problem hiding this comment.
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"oradvantage_normalization="token_globa"would silently select the token-KL or ordinary-whitening path;_assert_spec_implementations_resolveonly validates the three table identifiers. Validate the allowed values inAlgorithmSpecor 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_typeis accepted from YAML, butis_sftand the SFT/RL setup were computed before this helper runs. If an RL invocation overrides it tosft, 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. Rejectloss_typein 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_normalizedis implemented asreward_normalizer != "none", so this gate will treat every future non-nonenormalizer 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.
| ├── 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. |
| | 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 | |
| ├── rewards.py reward 归一化策略 + REWARD_NORMALIZERS | ||
| ├── advantages.py advantage 估计器 + ADVANTAGE_FNS | ||
| ├── policy.py policy loss 适配器 + POLICY_LOSS_FNS | ||
| └── numerics.py 共享的数值常量与退化判定 |
| | 参数 | 类型 | 默认值 | 可选值 | 说明 | | ||
| |------|------|--------|--------|------| | ||
| | `--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 系数启用 | |
| `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`). |
| _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.
a3ebe9d to
309cd81
Compare
|
两条阻塞都改完了,已 rebase 到 1. RLOO 挪进注册表了
这 11 条里有 5 条(禁 fully-async/hybrid、 main 带过来的 6 个 RLOO 测试文件我一个字没改( 2.
|
There was a problem hiding this comment.
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_batchappears inco_names, while the existing GAE tests callppo_utilsdirectly. A dropped/reorderedkl_coef, misplaced terminal reward, or omittedpadded_total_lengthswould therefore pass despite changing PPO (especially CP/VL) behavior. Add a test throughcompute_advantages_and_returnsthat compares both outputs with and without a non-Nonepadded layout.
return get_advantages_and_returns_batch(
total_lengths,
response_lengths,
values,
shaped_rewards,
| kl_level: str = "token" | ||
| """``"token"`` or ``"sequence"``; GSPO constrains the sequence as a whole.""" | ||
|
|
||
| advantage_normalization: str = "whiten" |
|
这轮相较上次进展很大:RLOO 已迁入 reward、advantage、policy loss 和参数能力注册;loss.py 两处重 不过,严格按照 Task 27 的验收要求,当前仍建议 Request changes。 1. [验收阻塞] critic/value 能力尚未真正由注册表驱动needs_critic 当前主要用于生成 ALGOS,但实际训练编排和 value 数据流仍在多处按 "ppo" 判断:
因此,注册第二个 needs_critic=True 的算法后,虽然 parser 和 ALGOS 会接受它,但 values 字段、 建议让这些调用点读取准确的 capability 或 handler identifier;如果部分行为并不等同于 2. [P1] custom YAML 中显式设置的 global batch 会被静默覆盖apply_custom_config_overrides() derive_global_batch_size(args, enforce_consistency=False) 只要 num_steps_per_rollout 非空,YAML 显式设置的 global_batch_size 就会先写入、随后被推导值静默 建议区分 YAML 是否显式包含 global_batch_size:
请补 num_steps_per_rollout != None 且 YAML 显式设置 global_batch_size 的回归测试。 3. [P2] PPO/GAE adapter 缙少执行级数值等价测试advantage_gae 题目明确要求现有算法数值等价性。建议通过统一 dispatcher 与 main 原逻辑逐位对拍,至少覆盖:
另外,kl_level 和 advantage_normalization 是有限枚举但运行时没有 fail-fast 校验,拼错会静默选择 完成这些后,#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.
|
三项都已处理,head 1. [验收阻塞] critic 能力由注册表驱动新增
没有复用 已补您要的测试:临时注册第二个 有两处 2. [P1] YAML global_batch_size 不再被静默覆盖
三条回归测试:冲突被拒、一致值存活、 3. [P2] GAE 数值等价与枚举校验对拍:替掉 一个覆盖边界要主动说明: 枚举校验:放在 |
【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 的 base98a1274上不成立:上游已经手工补进去了。这个例子来自更早的 main,写进 commit message 是我的疏忽,会诊时被指出。它恰恰说明手工维护的表会漂移,但不该被当成本 PR 基线上的现存 bug 引用。)AlgorithmSpec让这些事实只说一次。关键设计
字段存字符串标识符,不存 callable。 advantage 公式在
AdvantagesRay 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_lengthsloss.py:558(maybe_padded_total_lengths读args.qkv_format+ VL/unsplit-forward 标志)Advantagesdeployment 无法计算、不传(保持它原有行为)。这不是装饰性的:bshd / VL / unsplit-forward 下 GAE 要按 padded 偏移切 CP 分片,传None不报错,只是读错 token 位置normalize_advantages+ CP mask 重建(约 80 行)loss.py:620起loss.py。它在 advantage 算完之后跑,且依赖max_seq_lens与 CP 组——deployment 两者都没有loss.py:561not mpu.is_pipeline_last_stage()/components/advantages.py:152log_probs is None and values is None and rollout_log_probs is None顺带回答 review 里的一个问题:
max_seq_lens没有同类问题。get_advantages_and_returns_batch的签名里根本没有这个参数,它只服务于上表第二行的 normalize+CP 段。等价性论证
原 RFC 依赖「
ppo_utils.py与 main 逐字节相同」。这个前提已经失效:039ce87 → 98a1274之间该文件改了+474 −29。改用两条证据:get_reinforce_plus_plus_baseline_advantages上游已移除kl_coef形参(KL 改由独立 k2 loss 承担),旧 adapter 仍在传。其余 6 个逐参数一致,包括compute_policy_loss新增的eps_clip_c——它默认None,main 的loss.py也不传,adapter 行为与之相同。--disable-grpo-std-normalization门控),tests/algorithms/test_dispatch_parity_vs_main.py对每个 estimator 与ppo_utils的裸 kernel 做逐位比对(torch.equal,不用allclose——它的默认容差足以吞掉有偏/无偏标准差的差别)。Ask First 声明
按
CLAUDE.md,以下改动落在需要事先确认的范围内,在此明确列出请维护者裁定:relax/utils/arguments.py参数解析:--advantage-estimator的choices由硬编码列表改为list_algorithm_names()。取值集合本身不变(本 PR 不增不减算法)。relax/core/registry.py的ALGOS:由手写 dict 改为从注册表派生。process_role未改。这是 Controller 直接消费的编排面。validate_algorithm_args,插在_validate_reinforce_plus_plus_args之后——那个函数标着 frozen Task 29 contract 且其测试匹配特定文案,先跑会抢走它的错误消息。如果其中任何一条希望换个做法,我照改。
一处刻意不做
_validate_reinforce_plus_plus_args为reinforce_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
会把算法名分支恢复回来。所以是迁移进来,不是并排放着。
post_process_rewards里的== "rloo"内联分支group_leave_one_outnormalizerelif advantage_estimator == "rloo"的 policy lossPOLICY_LOSS_FNS["rloo"]"rloo"reward_normalizer/advantage_fn字段_compute_rloo_group_diagnostics的!= "rloo"门控reward_normalizer判定(指标 key 前缀保持rloo/不变,那是已发布的名字)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_normalization、forbids_normalize_advantages、min_group_size、forbids_reward_side_kl、requires_global_token_loss。其中forbids_reward_side_kl、requires_rewards_normalization、min_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_clip、requires_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 {放回去,测试立刻变红。测试
失败集合与
main@4899b8f在同一台机器上逐条一致(2 failed + 2 errors:/dev/shm在 macOS不存在,以及一个既有的
test_reward_router失败)。验证方式:在该 commit 上开一个 detachedworktree 重跑同样的用例,得到同一组失败。(早期版本说是用
git stash验证的——改动已经提交,stash 什么也不会动,那个说法是错的。)
pre-commit run --all-files全部通过。运行时证据(Modal,H100)
CPU 测试证明的是映射表和数值;Controller 真的据此把服务建起来、算法真的按注册表跑这一段
只能靠跑。四次冒烟,全部
succeeded:d11e247raysubmit_Lap4dJNZgGvFqyjw,8 步,loss 0.403/0.091/0.138/0.219d11e247raysubmit_TGe58MhMv62FZTb7,4 步,SeqlenBalancedSampler with dp_size=2d11e247raysubmit_9p3bUqYts8kkZ6ci,critic 服务 + 8 步 value_lossa61b68fraysubmit_HR4gXykqDnKtXC7X,4 步,pg_clipfrac恒为 0revision 归属,如实写出来而不是含糊成「都在 HEAD」:
PR2 的代码;对本 PR 有直接意义的是 PPO 与 RLOO 两次。
a61b68f,当前 PR2 tip 在relax/下只有一处可执行改动:
observed_reward_signal的except元组从(TypeError, ValueError)加宽为也含
KeyError, IndexError(仅指标路径)。arguments.py的非注释改动为零——YAML 校验补齐本来就在冒烟的那一版里,rebase 只是把它从 PR2 移到了本 PR。
PPO 是本 PR 改动最大的既有算法——唯一
needs_critic=True的那个,ALGOS里的 critic 条目从写死的名字比较改成由该字段推导:
RLOO 是本 PR 重写最多的那个(reward normalizer、policy loss adapter、11 条启动约束压成
6 个 capability 字段)。RLOO 本身在上游 #205(
98a7234)已有端到端记录——那条 commit message 写了 2-GPU canary 与 GRPO/RLOO 各 60/60 步的配对实验;这里跑的是注册化之后的这一版。三条独立证据说明跑的确实是 RLOO 而不是退回 PPO-clip:
启动校验一次误拦都没有。冒烟脚本从已知能跑的 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。这些靠读代码关不掉,只能靠跑。