From 309cd81912bbfbe6be43a5daaaac6bb3345000d2 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 13:11:21 +0800 Subject: [PATCH 1/6] refactor(algorithms): declare each algorithm once in a registry `--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 (#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. --- docs/.vitepress/config.mts | 2 + docs/en/guide/adding-an-algorithm.md | 177 ++++++ docs/en/guide/configuration.md | 2 +- docs/zh/guide/adding-an-algorithm.md | 127 +++++ docs/zh/guide/configuration.md | 2 +- .../post_process_genrm_swap.py | 35 +- relax/algorithms/__init__.py | 15 + relax/algorithms/advantages.py | 161 ++++++ relax/algorithms/policy.py | 75 +++ relax/algorithms/rewards.py | 138 +++++ relax/algorithms/spec.py | 226 ++++++++ relax/backends/megatron/loss.py | 142 ++--- relax/components/advantages.py | 73 +-- relax/core/registry.py | 93 ++-- relax/utils/arguments.py | 381 +++++++++---- relax/utils/metrics/metric_utils.py | 12 +- relax/utils/utils.py | 65 +-- tests/algorithms/__init__.py | 0 tests/algorithms/test_advantage_estimators.py | 107 ++++ tests/algorithms/test_algorithm_registry.py | 155 ++++++ tests/algorithms/test_algos_roles.py | 135 +++++ .../algorithms/test_arguments_spec_driven.py | 316 +++++++++++ .../test_dispatch_parity_vs_main.py | 524 ++++++++++++++++++ tests/algorithms/test_policy_loss_dispatch.py | 181 ++++++ .../test_post_process_rewards_dispatch.py | 169 ++++++ tests/algorithms/test_reward_normalizers.py | 194 +++++++ 26 files changed, 3131 insertions(+), 376 deletions(-) create mode 100644 docs/en/guide/adding-an-algorithm.md create mode 100644 docs/zh/guide/adding-an-algorithm.md create mode 100644 relax/algorithms/__init__.py create mode 100644 relax/algorithms/advantages.py create mode 100644 relax/algorithms/policy.py create mode 100644 relax/algorithms/rewards.py create mode 100644 relax/algorithms/spec.py create mode 100644 tests/algorithms/__init__.py create mode 100644 tests/algorithms/test_advantage_estimators.py create mode 100644 tests/algorithms/test_algorithm_registry.py create mode 100644 tests/algorithms/test_algos_roles.py create mode 100644 tests/algorithms/test_arguments_spec_driven.py create mode 100644 tests/algorithms/test_dispatch_parity_vs_main.py create mode 100644 tests/algorithms/test_policy_loss_dispatch.py create mode 100644 tests/algorithms/test_post_process_rewards_dispatch.py create mode 100644 tests/algorithms/test_reward_normalizers.py diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index c2157f054..b686efc29 100644 --- a/docs/.vitepress/config.mts +++ b/docs/.vitepress/config.mts @@ -243,6 +243,7 @@ export default defineConfig({ { text: 'Installation', link: '/en/guide/installation' }, { text: 'Quick Start', link: '/en/guide/quick-start' }, { text: 'Customize Training', link: '/en/guide/customize-training' }, + { text: 'Adding an Algorithm', link: '/en/guide/adding-an-algorithm' }, { text: 'SFT Training', link: '/en/guide/sft-training' }, { text: 'PPO Training', link: '/en/guide/ppo-training' }, { text: 'REINFORCE++', link: '/en/guide/reinforce-plus-plus' }, @@ -355,6 +356,7 @@ export default defineConfig({ { text: '安装', link: '/zh/guide/installation' }, { text: '快速上手', link: '/zh/guide/quick-start' }, { text: '自定义训练', link: '/zh/guide/customize-training' }, + { text: '接入新算法', link: '/zh/guide/adding-an-algorithm' }, { text: 'SFT 训练', link: '/zh/guide/sft-training' }, { text: 'PPO 训练', link: '/zh/guide/ppo-training' }, { text: 'REINFORCE++', link: '/zh/guide/reinforce-plus-plus' }, diff --git a/docs/en/guide/adding-an-algorithm.md b/docs/en/guide/adding-an-algorithm.md new file mode 100644 index 000000000..b415cf13d --- /dev/null +++ b/docs/en/guide/adding-an-algorithm.md @@ -0,0 +1,177 @@ +# Adding an Algorithm + +Algorithms plug into Relax through the registry under `relax/algorithms/`. An +algorithm name is no longer scattered across `if/elif` chains — it is described +by one `AlgorithmSpec`, and each stage looks up what it needs. + +## Registry Layout + +``` +relax/algorithms/ +├── spec.py AlgorithmSpec definition + the ALGORITHM_SPECS registry +├── rewards.py reward normalization strategies + REWARD_NORMALIZERS +├── advantages.py advantage estimators + ADVANTAGE_FNS +└── policy.py policy loss adapters + POLICY_LOSS_FNS +``` + +Three hard constraints: + +1. **No heavy top-level imports under `relax/algorithms/`** — not `megatron`, + `ray`, `transfer_queue`, `tensordict`, `relax.components` or + `relax.backends`. The registry is imported by argument parsing and by both + worker processes; one heavy import drags the whole training stack into + `--help` and into a CPU-only CI runner. Import inside the function when you + genuinely need one. +2. **Spec fields hold string identifiers, not callables.** The advantage + computation runs in the Ray Serve `Advantages` process while the policy loss + runs in the Megatron worker, and those two import different module subsets. + Only the algorithm name crosses the process boundary; each side resolves it + against its own table. +3. **Do not hand-edit the `ALGOS` role table.** It is derived from the registry, + so a new algorithm gets the standard RL role set automatically. + +## How Much Does Adding One Cost + +Honestly: **not "one dict entry".** + +| Situation | Files to touch | +|---|---| +| Reuses existing reward normalization / advantage / policy loss, just combined differently | 1 (`spec.py`) | +| Needs new maths (a new advantage formula, say) | 2-3 (`spec.py` plus the implementation module) | +| Also needs new command-line options | 4-6 (the above, plus the option and its validation in `arguments.py`, plus an example and docs) | + +What the registry removes is one algorithm name being interpreted in six +scattered if/elif chains — not the cost of adding an algorithm. An algorithm +that needs both new maths and new options lands in the last row. + +The `ALGOS` role table is the one part that genuinely costs nothing: it derives +itself from the registry. + +## Steps + +### 1. Add a spec entry + +Edit `ALGORITHM_SPECS` in `relax/algorithms/spec.py`: + +```python +"my_algo": AlgorithmSpec( + name="my_algo", + reward_normalizer="group_mean_std", # reuse an existing one, or see step 2 + advantage_fn="grpo_broadcast", + policy_loss_fn="ppo_clip", +), +``` + +If your algorithm is identical to an existing one at some stage, reuse that +identifier. GRPO, GSPO, SAPO and CISPO are equivalent at the advantage layer, +so all four share `"grpo_broadcast"`. + +Capability fields: + +| Field | Effect | +|-------|--------| +| `kl_level` | `"token"` or `"sequence"` (GSPO constrains the sequence) | +| `needs_full_log_probs` | Whether the loss needs CP-gathered full log probs | +| `advantage_normalization` | What `--normalize-advantages` does: `"whiten"` (masked whitening) or `"token_global"` (REINFORCE++'s global token-level normalization, which also switches on the mask-safe loss reducer) | +| `needs_critic` | Whether a critic service is required; drives `args.use_critic` | +| `requires_normalize_advantages` | Demand `--normalize-advantages` | +| `forbids_normalize_advantages` | Reject `--normalize-advantages` (the estimator keeps the advantage's scale on purpose) | +| `requires_rewards_normalization` | Reject `--disable-rewards-normalization` | +| `min_group_size` | Floor on `--n-samples-per-prompt` | +| `forbids_reward_side_kl` | Demand `--kl-coef 0`; there is nowhere to put a reward-side KL term (`--use-kl-loss` is unaffected) | +| `requires_global_token_loss` | Demand `--calculate-per-token-loss`; the per-sample token-mean reducer would reweight responses by `1 / response_length` | +| `requires_on_policy_updates` | Rejects five knobs at once: `--fully-async` / `--hybrid`, `--max-staleness != 0`, `--num-steps-per-rollout != 1`, `rollout_batch_size * n_samples != global_batch_size`, and `--partial-rollout` / `--use-dynamic-global-batch-size`. For objectives with no importance-ratio correction | + +The four `validate_*` functions in `relax/utils/arguments.py` consume every +field in that table except `kl_level`, `needs_full_log_probs` and +`advantage_normalization`, so for the rest, declaring the field is enough — you +do not add an `if` there. (They are four rather than one because argument +validation has a derivation order: `--kl-coef` has to be settled before +validation demands that `--ref-load` exist on disk, and the one-update equality +cannot be checked until `global_batch_size` has taken its final value. Neither +has anything to do with the algorithm being special.) Those three fields are +read in `relax/backends/megatron/loss.py` instead: a genuinely new value needs a +branch there, an existing one does not. + +### 2. Write pure functions for genuinely new maths + +Only needed when your algorithm differs from every existing one at that stage. + +**Reward normalization** (`relax/algorithms/rewards.py`), signature +`fn(args, samples, raw_rewards) -> list[float]`: + +```python +def normalize_my_strategy(args, samples, raw_rewards): + positions_by_group = group_positions(samples, args.n_samples_per_prompt) + ... + return normalized # one scalar per sample + +REWARD_NORMALIZERS["my_strategy"] = normalize_my_strategy +``` + +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. + +**Advantage estimator** (`relax/algorithms/advantages.py`), signature +`fn(args, *, rewards, kl, loss_masks, response_lengths, total_lengths, values)` +returning `(advantages, returns)`, both `list[Tensor]`: + +```python +def advantage_my_algo(args, *, rewards, kl, **_unused): + ... + return advantages, returns + +ADVANTAGE_FNS["my_algo"] = advantage_my_algo +``` + +**Policy loss** (`relax/algorithms/policy.py`), signature +`fn(args, *, log_probs, ppo_kl, advantages) -> (pg_loss, pg_clipfrac)`. The +underlying kernels take different argument lists; the adapter normalizes them. + +### 3. Write unit tests + +Tests under `tests/algorithms/` need only torch — no megatron, ray or +transfer_queue: + +```bash +pytest tests/algorithms/ -v +``` + +Cover at least: + +- Registration and dispatch: the name is in `ALGORITHM_SPECS`, capability fields + match expectations, an unregistered name raises. +- Numerics: hand-compute a small case as the reference. Do not use all-zero or + all-equal rewards — every formula returns 0 on those, so the test proves + nothing. +- Degenerate cases: a group where all rewards are equal, boundary values of + `n_samples_per_prompt`, missing fields, non-numeric input. +- **When changing an existing algorithm**: freeze the old implementation into + the test file as a reference and compare bit-for-bit + (`view(torch.int32).equal`). Do not use `allclose` — its default tolerance is + wide enough to swallow the difference between a biased and an unbiased + standard deviation. `tests/algorithms/test_reward_normalizers.py` is a + worked example. + +### 4. Add an example and documentation + +- `examples//`: a launch script, plus a custom reward function if needed. +- `docs/{zh,en}/examples/algorithms.md`: how it works, the parameter table, a + quick start, and **known deviations** — write down where the implementation + differs from the paper rather than leaving users to discover it. + +## Arguments + +Algorithm-specific options go in `add_algo_arguments` in +`relax/utils/arguments.py`. The `--advantage-estimator` choices come from +`list_algorithm_names()`, so registering is enough; there is no name list to +maintain. + +Put cross-argument validation in `validate_algorithm_args`, and prefer +expressing it through a spec field over comparing algorithm names — the latter +is exactly what this registry exists to remove. + +## References + +- [Algorithm Reference](../examples/algorithms.md) diff --git a/docs/en/guide/configuration.md b/docs/en/guide/configuration.md index f98a14739..0f8e54028 100644 --- a/docs/en/guide/configuration.md +++ b/docs/en/guide/configuration.md @@ -299,7 +299,7 @@ bash scripts/training/text/run-qwen3-4B-fp16-8xgpu.sh \ | 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`, `rloo`, `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 | | `--normalize-advantages` | flag | False | - | Whether to normalize advantages | | `--disable-grpo-std-normalization` | flag | - | - | Disable GRPO standard deviation normalization (from [Dr.GRPO](https://arxiv.org/pdf/2503.20783)) | | `--disable-rewards-normalization` | flag | - | - | Disable reward normalization | diff --git a/docs/zh/guide/adding-an-algorithm.md b/docs/zh/guide/adding-an-algorithm.md new file mode 100644 index 000000000..ebeb5f9d8 --- /dev/null +++ b/docs/zh/guide/adding-an-algorithm.md @@ -0,0 +1,127 @@ +# 接入一个新算法 + +Relax 的算法通过 `relax/algorithms/` 下的注册表接入。一个算法名不再散落在各处的 `if/elif` 里——它由一条 `AlgorithmSpec` 描述，各个阶段按需查表。 + +## 注册表的结构 + +``` +relax/algorithms/ +├── spec.py AlgorithmSpec 定义 + ALGORITHM_SPECS 注册表 +├── rewards.py reward 归一化策略 + REWARD_NORMALIZERS +├── advantages.py advantage 估计器 + ADVANTAGE_FNS +└── policy.py policy loss 适配器 + POLICY_LOSS_FNS +``` + +三条硬约束： + +1. **`relax/algorithms/` 下禁止顶层 import 重依赖**——不能有 `megatron`、`ray`、`transfer_queue`、`tensordict`、`relax.components`、`relax.backends`。注册表会被参数解析和两个 worker 进程 import；一个重依赖会把整个训练栈拖进 `--help` 和只有 CPU 的 CI。确实需要时在函数内 import。 +2. **spec 的字段存字符串标识符，不存函数引用**。advantage 计算跑在 Ray Serve 的 `Advantages` 进程，policy loss 跑在 Megatron worker 进程，两者 import 的模块子集不同。跨进程只传算法名，各进程本地查表。 +3. **`ALGOS` 角色表不用手改**。它从注册表自动派生，新算法自动获得标准 RL 角色集合。 + +## 接入一个新算法要改多少 + +先说实话：**不是「加一条 dict entry」就完事**。 + +| 情况 | 要改的文件 | +|---|---| +| 复用现成的 reward 归一化 / advantage / policy loss，只是组合方式不同 | 1 个（`spec.py`） | +| 需要一种新的数学（如新的 advantage 公式） | 2–3 个（`spec.py` + 对应的实现模块） | +| 还需要新的命令行参数 | 4–6 个（上述 + `arguments.py` 的参数声明与校验 + 示例 + 文档） | + +注册表消除的是「同一个算法名散落在 6 处 if/elif」，不是「新增算法零成本」。 + +`ALGOS` 角色表是唯一真正做到零改动的部分——它从注册表自动派生。 + +## 步骤 + +### 1. 加一条 spec + +编辑 `relax/algorithms/spec.py` 的 `ALGORITHM_SPECS`： + +```python +"my_algo": AlgorithmSpec( + name="my_algo", + reward_normalizer="group_mean_std", # 复用现成的，或见第 2 步 + advantage_fn="grpo_broadcast", + policy_loss_fn="ppo_clip", +), +``` + +如果新算法在某个阶段与已有算法完全一致，直接复用那个标识符即可——例如 GRPO / GSPO / SAPO / CISPO 在 advantage 层完全等价，四者共享 `"grpo_broadcast"`。 + +可用的能力字段： + +| 字段 | 作用 | +|------|------| +| `kl_level` | `"token"` 或 `"sequence"`（GSPO 用序列级） | +| `needs_full_log_probs` | loss 是否需要 CP all-gather 后的完整 log probs | +| `advantage_normalization` | `--normalize-advantages` 的归一化方式：`"whiten"`（掩码白化）或 `"token_global"`（REINFORCE++ 的全局 token 级归一化，同时切换掩码安全的 loss reducer） | +| `needs_critic` | 是否需要 critic 服务，驱动 `args.use_critic` | +| `requires_normalize_advantages` | 强制要求 `--normalize-advantages` | +| `forbids_normalize_advantages` | 禁止 `--normalize-advantages`（算法刻意保留了 advantage 的尺度时） | +| `requires_rewards_normalization` | 禁止 `--disable-rewards-normalization` | +| `min_group_size` | `--n-samples-per-prompt` 的下限 | +| `forbids_reward_side_kl` | 要求 `--kl-coef 0`（reward 侧 KL 项无处可放；`--use-kl-loss` 不受影响） | +| `requires_global_token_loss` | 强制要求 `--calculate-per-token-loss`（否则按样本取 token 均值，会按 `1 / response_length` 重新加权） | +| `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`。适用于没有重要性比值修正的目标函数 | + +表里除 `kl_level`、`needs_full_log_probs` 和 `advantage_normalization` 之外的字段，都由 `relax/utils/arguments.py` 的四个 `validate_*` 函数统一消费，**声明即生效**，不需要再去 `arguments.py` 加 `if`。（拆成四个是因为参数校验本身有推导顺序——例如 `--kl-coef` 必须在「检查 `--ref-load` 是否存在」之前判掉，one-update 等式必须在 `global_batch_size` 定稿之后判——与算法特殊性无关。）那三个字段是在 `relax/backends/megatron/loss.py` 里读的：新增一个前所未有的取值需要在那里加分支，复用已有取值则不用。 + +### 2. 需要新公式时，写纯函数并登记 + +只有当新算法在某个阶段的数学与现有算法都不同时才需要这一步。 + +**Reward 归一化**（`relax/algorithms/rewards.py`），签名固定为 `fn(args, samples, raw_rewards) -> list[float]`： + +```python +def normalize_my_strategy(args, samples, raw_rewards): + positions_by_group = group_positions(samples, args.n_samples_per_prompt) + ... + return normalized # 每个 sample 一个标量 + +REWARD_NORMALIZERS["my_strategy"] = normalize_my_strategy +``` + +产出必须是**每个 sample 一个标量**。这条约束让 TransferQueue 的 schema 保持不变——即使算法内部要看多个奖励分量，也要在这一层收敛成一个标量。 + +**Advantage 估计器**（`relax/algorithms/advantages.py`），签名 `fn(args, *, rewards, kl, loss_masks, response_lengths, total_lengths, values) -> (advantages, returns)`，两者都是 `list[Tensor]`： + +```python +def advantage_my_algo(args, *, rewards, kl, **_unused): + ... + return advantages, returns + +ADVANTAGE_FNS["my_algo"] = advantage_my_algo +``` + +**Policy loss**（`relax/algorithms/policy.py`），签名 `fn(args, *, log_probs, ppo_kl, advantages) -> (pg_loss, pg_clipfrac)`。底层算子签名不一致，适配器负责统一。 + +### 3. 写单测 + +`tests/algorithms/` 下的测试不依赖 megatron / ray / transfer_queue，只要 torch 就能跑： + +```bash +pytest tests/algorithms/ -v +``` + +至少覆盖： + +- 注册与分发：算法名在 `ALGORITHM_SPECS` 里；能力字段与预期一致；未注册名报错。 +- 数值：手算一个小例子做对照，别用全零或全相同的 reward——那种输入下任何公式都输出 0，测不出东西。 +- 退化场景：组内 reward 全相同、`n_samples_per_prompt` 取边界值、缺字段、非数值输入。 +- **改动已有算法时**：把旧实现冻结进测试文件当参照，逐位对拍（`view(torch.int32).equal`），不要用 `allclose`——它的默认容差足以吞掉无偏/有偏标准差的差异。`tests/algorithms/test_reward_normalizers.py` 是现成范例。 + +### 4. 加示例与文档 + +- `examples//`：启动脚本，必要时附自定义 reward 函数。 +- `docs/{zh,en}/examples/algorithms.md`：算法原理、关键参数表、快速开始，以及**已知偏差**——实现与论文不一致的地方要写出来，不要留给使用者去发现。 + +## 参数 + +新增算法专用参数时改 `relax/utils/arguments.py` 的 `add_algo_arguments`。`--advantage-estimator` 的 `choices` 由 `list_algorithm_names()` 生成，注册即可用，不需要手动维护名单。 + +跨参数的校验写进 `validate_algorithm_args`，并优先用 spec 字段表达而不是比较算法名——后者正是这套注册表要消除的东西。 + +## 参考 + +- [算法参考](../examples/algorithms.md) diff --git a/docs/zh/guide/configuration.md b/docs/zh/guide/configuration.md index 9da9012ac..957029903 100644 --- a/docs/zh/guide/configuration.md +++ b/docs/zh/guide/configuration.md @@ -299,7 +299,7 @@ bash scripts/training/text/run-qwen3-4B-fp16-8xgpu.sh \ | 参数 | 类型 | 默认值 | 可选值 | 说明 | |------|------|--------|--------|------| -| `--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`、`rloo`、`ppo`、`reinforce_plus_plus`、`reinforce_plus_plus_baseline` | 优势估计器。以 `--help` 为准：取值直接读注册表，新增算法无需改这张表即可出现。OPD 独立于该选项，使用 `--use-opd` 及对应 KL/loss 系数启用 | | `--normalize-advantages` | flag | False | - | 是否归一化优势 | | `--disable-grpo-std-normalization` | flag | - | - | 禁用 GRPO 标准差归一化（来自 [Dr.GRPO](https://arxiv.org/pdf/2503.20783)） | | `--disable-rewards-normalization` | flag | - | - | 禁用 reward 归一化 | diff --git a/examples/generate_reward_model/post_process_genrm_swap.py b/examples/generate_reward_model/post_process_genrm_swap.py index b3dcf4697..f3556b99d 100644 --- a/examples/generate_reward_model/post_process_genrm_swap.py +++ b/examples/generate_reward_model/post_process_genrm_swap.py @@ -100,19 +100,26 @@ async def _score_all(samples): def _grpo_normalize(args, raw_rewards): - """Replicates the default GRPO group normalization in - relax.utils.utils.post_process_rewards.""" - if ( - args.advantage_estimator - not in [ - "grpo", - "gspo", - "sapo", - "cispo", - "reinforce_plus_plus_baseline", - ] - or not args.rewards_normalization - ): + """Replicates the default group normalization in + relax.utils.utils.post_process_rewards. + + Which algorithms normalize, and which of those also divide by the group + standard deviation, comes from the algorithm registry rather than from a + copy of the whitelist — a copy would silently go stale the next time an + algorithm is added. + + The gate names the two normalizers this function actually reimplements + rather than asking ``spec.is_group_normalized``: that property would also be + true of a future normalizer computing something else entirely, and this + reimplementation would then silently diverge from it. Naming normalizers + keeps it registry-driven -- a new algorithm reusing either one is covered + for free, and a genuinely new normalizer is exactly the case where a human + needs to look at this function. + """ + from relax.algorithms import get_algorithm + + spec = get_algorithm(args.advantage_estimator) + if spec.reward_normalizer not in ("group_mean", "group_mean_std") or not args.rewards_normalization: return raw_rewards rewards = torch.tensor(raw_rewards, dtype=torch.float) if rewards.shape[-1] == args.n_samples_per_prompt * args.rollout_batch_size: @@ -121,7 +128,7 @@ def _grpo_normalize(args, raw_rewards): rewards = rewards.view(-1, rewards.shape[-1]) mean = rewards.mean(dim=-1, keepdim=True) rewards = rewards - mean - if args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo"] and args.grpo_std_normalization: + if spec.reward_normalizer == "group_mean_std" and args.grpo_std_normalization: std = rewards.std(dim=-1, keepdim=True) rewards = rewards / (std + 1e-6) return rewards.flatten().tolist() diff --git a/relax/algorithms/__init__.py b/relax/algorithms/__init__.py new file mode 100644 index 000000000..2d8c798bf --- /dev/null +++ b/relax/algorithms/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Algorithm registry: names, capabilities and implementations in one place. + +Nothing in this package may import ``megatron``, ``ray``, ``transfer_queue``, +``tensordict``, ``relax.components`` or ``relax.backends`` at module level. The +registry is imported by argument parsing and by both worker processes, so a +heavy import here would pull the whole training stack into ``--help`` and into +the CPU-only test runner. +""" + +from relax.algorithms.spec import ALGORITHM_SPECS, AlgorithmSpec, get_algorithm, list_algorithm_names + + +__all__ = ["ALGORITHM_SPECS", "AlgorithmSpec", "get_algorithm", "list_algorithm_names"] diff --git a/relax/algorithms/advantages.py b/relax/algorithms/advantages.py new file mode 100644 index 000000000..a93129ad6 --- /dev/null +++ b/relax/algorithms/advantages.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Advantage estimators as pure functions shared by both execution paths. + +The colocate path calls these from ``relax.backends.megatron.loss`` inside the +Megatron worker; the fully-async path calls them from the +``relax.components.advantages`` Ray Serve deployment. Keeping the maths here +means the two call sites differ only in what surrounds them: the pipeline-stage +early return, in-place write-back versus nested-tensor packing, and the +optional advantage whitening. + +Note that the group-wise reward standardisation happened earlier, on the +rollout side (see :mod:`relax.algorithms.rewards`). By the time an estimator +runs, ``rewards`` already holds one normalised scalar per sample. +""" + +from typing import Any, Callable + +import torch + +from relax.algorithms.spec import get_algorithm +from relax.utils.training.ppo_utils import ( + get_advantages_and_returns_batch, + get_grpo_returns, + get_reinforce_plus_plus_baseline_advantages, + get_reinforce_plus_plus_returns, +) + + +def _as_reward_tensor(rewards: Any, kl: list[torch.Tensor]) -> torch.Tensor: + if isinstance(rewards, torch.Tensor): + return rewards.to(dtype=torch.float32, device=kl[0].device) + return torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) + + +def advantage_grpo_broadcast(args: Any, *, rewards, kl, **_unused): + """Broadcast the already group-normalised scalar reward over tokens.""" + reward_tensor = _as_reward_tensor(rewards, kl) + returns = get_grpo_returns(reward_tensor, kl) + advantages = list(returns) # separate list so rebinding one does not move the other + return advantages, returns + + +def advantage_reinforce_plus_plus(args: Any, *, rewards, kl, loss_masks, response_lengths, total_lengths, **_unused): + """Discounted returns for REINFORCE++ + (https://arxiv.org/pdf/2501.03262).""" + reward_tensor = _as_reward_tensor(rewards, kl) + returns = get_reinforce_plus_plus_returns( + rewards=reward_tensor, + kl=kl, + loss_masks=loss_masks, + response_lengths=response_lengths, + total_lengths=total_lengths, + kl_coef=args.kl_coef, + gamma=args.gamma, + ) + advantages = list(returns) + return advantages, returns + + +def advantage_reinforce_plus_plus_baseline(args: Any, *, rewards, kl, loss_masks, **_unused): + """REINFORCE++ with a group baseline already subtracted upstream. + + No ``kl_coef``: this estimator keeps the KL penalty out of the advantage + (``_validate_reinforce_plus_plus_args`` requires ``--kl-coef 0`` and an + independent k2 loss instead), and the helper dropped the parameter to + match. + """ + reward_tensor = _as_reward_tensor(rewards, kl) + advantages = get_reinforce_plus_plus_baseline_advantages( + rewards=reward_tensor, + kl=kl, + loss_masks=loss_masks, + ) + # NOTE(dev): returns aliases the same list object, matching the pre-refactor + # behaviour. On-policy distillation rebinds slots of the advantages list and + # both views are expected to observe that. + return advantages, advantages + + +def advantage_gae( + args: Any, + *, + rewards, + kl, + values, + response_lengths, + total_lengths, + padded_total_lengths=None, + **_unused, +): + """Generalised advantage estimation (PPO). + + ``padded_total_lengths`` is what the two call sites disagree on, and it is + load-bearing rather than cosmetic: under bshd ``qkv_format``, VL models or + unsplit forward, the sequences were padded before the forward pass, so CP + un-sharding inside ``get_advantages_and_returns_batch`` has to slice at the + padded offsets. Passing ``None`` there does not raise — it silently reads + the wrong token positions. The Megatron path computes the value via + ``maybe_padded_total_lengths`` and passes it; the ``Advantages`` deployment + has no equivalent and passes ``None``, which is the behaviour it had before + this handler existed. + """ + from megatron.core import mpu + + shaped_rewards = [] + cp_rank = mpu.get_context_parallel_rank() + for reward, k in zip(rewards, kl, strict=False): + k *= -args.kl_coef + if cp_rank == 0: + k[-1] += reward + shaped_rewards.append(k) + return get_advantages_and_returns_batch( + total_lengths, + response_lengths, + values, + shaped_rewards, + args.gamma, + args.lambd, + padded_total_lengths=padded_total_lengths, + ) + + +ADVANTAGE_FNS: dict[str, Callable[..., tuple[list[torch.Tensor], list[torch.Tensor]]]] = { + "grpo_broadcast": advantage_grpo_broadcast, + "reinforce_plus_plus": advantage_reinforce_plus_plus, + "reinforce_plus_plus_baseline": advantage_reinforce_plus_plus_baseline, + "gae": advantage_gae, +} + + +def compute_advantages_and_returns( + args: Any, + *, + rewards, + kl: list[torch.Tensor], + loss_masks: list[torch.Tensor] | None = None, + response_lengths: list[int] | None = None, + total_lengths: list[int] | None = None, + values: list[torch.Tensor] | None = None, + padded_total_lengths: list[int] | None = None, +) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Dispatch to the estimator registered for ``args.advantage_estimator``. + + ``padded_total_lengths`` is likewise call-site specific: only the Megatron + path can compute it, and only GAE consumes it. Every parameter here is + keyword-only and every estimator absorbs the rest, so a call site that + cannot supply one omits it rather than inventing a value. + """ + spec = get_algorithm(args.advantage_estimator) + fn = ADVANTAGE_FNS[spec.advantage_fn] + return fn( + args, + rewards=rewards, + kl=kl, + loss_masks=loss_masks, + response_lengths=response_lengths, + total_lengths=total_lengths, + values=values, + padded_total_lengths=padded_total_lengths, + ) diff --git a/relax/algorithms/policy.py b/relax/algorithms/policy.py new file mode 100644 index 000000000..cc033c9fa --- /dev/null +++ b/relax/algorithms/policy.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Policy loss variants behind one uniform signature. + +The kernels in :mod:`relax.utils.training.ppo_utils` take different argument +lists. These adapters normalise them to ``fn(args, *, log_probs, ppo_kl, +advantages)`` so the caller can look one up by name instead of branching on the +algorithm. Adding a variant means adding an adapter and one registry entry; no +call site changes. +""" + +from typing import Any, Callable + +import torch + +from relax.algorithms.spec import get_algorithm +from relax.utils.training.ppo_utils import ( + compute_cispo_loss, + compute_policy_loss, + compute_rloo_loss, + compute_sapo_loss, +) + + +def policy_loss_ppo_clip(args: Any, *, log_probs, ppo_kl, advantages): + """Standard clipped surrogate objective (GRPO, GSPO, PPO, REINFORCE++).""" + return compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) + + +def policy_loss_sapo(args: Any, *, log_probs, ppo_kl, advantages): + """Smooth trust region: sigmoid gating instead of a hard clip.""" + return compute_sapo_loss( + ppo_kl=ppo_kl, + advantages=advantages, + tau_pos=getattr(args, "sapo_tau_pos", 1.0), + tau_neg=getattr(args, "sapo_tau_neg", 1.05), + ) + + +def policy_loss_cispo(args: Any, *, log_probs, ppo_kl, advantages): + """Clipped importance ratio that preserves the gradient direction.""" + return compute_cispo_loss( + log_probs=log_probs, + ppo_kl=ppo_kl, + advantages=advantages, + eps_clip=args.eps_clip, + eps_clip_high=args.eps_clip_high, + ) + + +def policy_loss_rloo(args: Any, *, log_probs, ppo_kl, advantages): + """Unclipped REINFORCE objective: ``-stopgrad(A) * log pi(y)``. + + ``ppo_kl`` is accepted and deliberately unused. Every other variant here + corrects for the policy having moved since the rollout; this one has no + such term, which is exactly why ``rloo`` declares + ``requires_on_policy_updates`` -- the correction is missing from the maths, + so it has to be guaranteed by the configuration instead. + """ + return compute_rloo_loss(log_probs=log_probs, advantages=advantages) + + +POLICY_LOSS_FNS: dict[str, Callable[..., tuple[torch.Tensor, torch.Tensor]]] = { + "ppo_clip": policy_loss_ppo_clip, + "sapo": policy_loss_sapo, + "cispo": policy_loss_cispo, + "rloo": policy_loss_rloo, +} + + +def compute_policy_loss_for(args: Any, *, log_probs, ppo_kl, advantages): + """Dispatch to the policy loss registered for + ``args.advantage_estimator``.""" + spec = get_algorithm(args.advantage_estimator) + return POLICY_LOSS_FNS[spec.policy_loss_fn](args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py new file mode 100644 index 000000000..9cc7efd22 --- /dev/null +++ b/relax/algorithms/rewards.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Reward normalisation strategies, one per algorithm family. + +These run on the rollout side (CPU) from +``relax.utils.utils.post_process_rewards``. Every normaliser takes the raw +per-sample scalar rewards and returns the per-sample scalars written into the +TransferQueue ``rewards`` column, so adding a strategy never changes the data +schema. +""" + +from typing import Any, Callable + +import torch + +from relax.utils.training.ppo_utils import compute_rloo_leave_one_out_rewards + + +GROUP_EPS = 1e-6 +"""Epsilon added to a group standard deviation before dividing by it. + +Matches the value the pre-registry GRPO path used, so GRPO, GSPO, SAPO and +CISPO keep producing exactly the numbers they produced before the registry +existed. That parity is what the equivalence tests assert, so this constant is +not free to move. +""" + + +def group_positions(samples: list[Any], expected_size: int) -> dict[int, list[int]]: + """Map ``Sample.group_index`` to the positions it occupies in ``samples``. + + Grouping follows ``group_index`` rather than position, so how the caller + orders the batch does not affect the result. + """ + positions_by_group: dict[int, list[int]] = {} + for position, sample in enumerate(samples): + if sample.group_index is None: + raise ValueError("Sample.group_index is required for group reward normalization.") + if sample.group_index not in positions_by_group: + positions_by_group[sample.group_index] = [] + positions_by_group[sample.group_index].append(position) + + for group_index, positions in positions_by_group.items(): + if len(positions) != expected_size: + raise ValueError(f"Reward group {group_index} has {len(positions)} samples, expected {expected_size}.") + return positions_by_group + + +def _group_normalize(args: Any, samples: list[Any], raw_rewards: list[float], *, use_std: bool) -> list[float]: + rewards = torch.tensor(raw_rewards, dtype=torch.float) + positions_by_group = group_positions(samples, args.n_samples_per_prompt) + + normalized_rewards = torch.empty_like(rewards) + for positions in positions_by_group.values(): + group_rewards = rewards[positions] + group_rewards = group_rewards - group_rewards.mean() + if use_std: + group_rewards = group_rewards / (group_rewards.std() + GROUP_EPS) + normalized_rewards[positions] = group_rewards + + return normalized_rewards.tolist() + + +def normalize_none(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: + """No normalisation — the estimator consumes raw rewards (REINFORCE++, + PPO).""" + return raw_rewards + + +def normalize_group_mean(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: + """Subtract the group mean only (REINFORCE++ baseline).""" + return _group_normalize(args, samples, raw_rewards, use_std=False) + + +def normalize_group_mean_std(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: + """Subtract the group mean, then optionally divide by the group std. + + ``--disable-grpo-std-normalization`` (Dr.GRPO) turns the division off. + """ + return _group_normalize(args, samples, raw_rewards, use_std=args.grpo_std_normalization) + + +def _reject_non_finite_group(group_index: int, positions: list[int], group_rewards: torch.Tensor) -> None: + """Raise if any reward in one group is NaN or infinite, naming the sample. + + The other normalisers do not check: an infinite reward there poisons only + the sample that carries it (``group_mean_std`` divides it away, and the + caller sees one bad advantage). A leave-one-out baseline averages the + *other* samples, so one bad value silently propagates into every advantage + in the group -- which is why this reports the offender's position instead + of letting the arithmetic swallow it. + """ + finite_mask = torch.isfinite(group_rewards) + if finite_mask.all(): + return + invalid_group_positions = (~finite_mask).nonzero(as_tuple=False).flatten().tolist() + invalid_sample_positions = [positions[position] for position in invalid_group_positions] + invalid_values = group_rewards[~finite_mask].tolist() + raise ValueError( + f"RLOO group_index={group_index} contains non-finite reward(s) at " + f"group position(s) {invalid_group_positions}, sample position(s) " + f"{invalid_sample_positions}: {invalid_values}." + ) + + +def normalize_group_leave_one_out(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: + """RLOO's leave-one-out baseline (Ahmadian et al. 2024, arXiv:2402.14740). + + Each sample is centred on the mean of the *others* in its prompt group:: + + A_i = r_i - mean_{j != i}(r_j) = G / (G - 1) * (r_i - mean(r)) + + Unlike :func:`normalize_group_mean_std` there is no division by the group + standard deviation, so the advantage keeps the reward's scale. The + right-hand identity is what + :func:`~relax.utils.training.ppo_utils.compute_rloo_leave_one_out_rewards` + computes; that helper also owns the ``G >= 2`` and finiteness contracts, + and it is shared with the rollout diagnostics so the numbers reported and + the numbers trained on cannot drift apart. + """ + rewards = torch.tensor(raw_rewards, dtype=torch.float) + positions_by_group = group_positions(samples, args.n_samples_per_prompt) + + normalized_rewards = torch.empty_like(rewards) + for group_index, positions in positions_by_group.items(): + group_rewards = rewards[positions] + _reject_non_finite_group(group_index, positions, group_rewards) + normalized_rewards[positions] = compute_rloo_leave_one_out_rewards(group_rewards) + + return normalized_rewards.tolist() + + +REWARD_NORMALIZERS: dict[str, Callable[[Any, list[Any], list[float]], list[float]]] = { + "none": normalize_none, + "group_mean": normalize_group_mean, + "group_mean_std": normalize_group_mean_std, + "group_leave_one_out": normalize_group_leave_one_out, +} diff --git a/relax/algorithms/spec.py b/relax/algorithms/spec.py new file mode 100644 index 000000000..d124845db --- /dev/null +++ b/relax/algorithms/spec.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Declarative descriptions of the RL algorithms Relax supports. + +A single ``--advantage-estimator`` value used to be interpreted independently by +six different files: role lookup, reward normalisation, the advantage formula, +the policy loss formula, and two rounds of argument validation. Adding an +algorithm meant finding all of them. ``AlgorithmSpec`` collects that metadata +in one place, so a new algorithm is one dict entry plus its pure functions. + +Fields hold **string identifiers**, not callables. The advantage formula runs +inside the Ray Serve ``Advantages`` deployment while the policy loss runs inside +the Megatron worker; those two processes import different module subsets, so we +ship the algorithm name and let each process resolve the identifier against its +own table. That also keeps this module free of heavy imports, which is what +makes the registry testable on a CPU-only runner. +""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AlgorithmSpec: + """Everything the training pipeline needs to know about one algorithm.""" + + name: str + + # --- reward stage (rollout side, CPU, scalar in / scalar out) --- + reward_normalizer: str + """Key into :data:`relax.algorithms.rewards.REWARD_NORMALIZERS`.""" + + # --- advantage stage --- + advantage_fn: str + """Key into :data:`relax.algorithms.advantages.ADVANTAGE_FNS`.""" + + # --- policy loss stage --- + policy_loss_fn: str + """Key into :data:`relax.algorithms.policy.POLICY_LOSS_FNS`.""" + + kl_level: str = "token" + """``"token"`` or ``"sequence"``; GSPO constrains the sequence as a whole.""" + + advantage_normalization: str = "whiten" + """How ``--normalize-advantages`` normalises, read in + ``relax/backends/megatron/loss.py``. + + ``"whiten"`` is the masked whitening every algorithm used before REINFORCE++ + arrived. ``"token_global"`` is REINFORCE++'s global token-level + normalisation, which also requires the mask-safe loss reducer -- the two go + together, which is why one field drives both call sites rather than two. + + Those two call sites were the last place a *maths* decision was still made + by comparing algorithm names, and unlike the remaining ``== "ppo"`` checks + this one is not expressible through an existing field: it is orthogonal to + ``needs_critic``. + """ + + needs_full_log_probs: bool = False + """Whether the loss needs CP-gathered full-response log probs.""" + + # --- orchestration and validation --- + needs_critic: bool = False + requires_normalize_advantages: bool = False + forbids_normalize_advantages: bool = False + """Set when the token-level ``--normalize-advantages`` pass would undo what + the estimator deliberately kept. + + RLOO's advantage carries the reward scale on purpose; re-whitening it after + DP sharding puts back the standard-deviation division RLOO removes, and + makes the result depend on how the batch was partitioned. + """ + + requires_rewards_normalization: bool = False + """Set when the algorithm's reward stage is load-bearing rather than + optional, so ``--disable-rewards-normalization`` would silently skip it.""" + + min_group_size: int = 1 + """Smallest ``--n-samples-per-prompt`` the reward stage is defined for.""" + + forbids_reward_side_kl: bool = False + """Set when the estimator has no place to put a reward-side KL penalty. + + Both members here compute their advantage from a completion-level signal + that the ``--kl-coef`` shaping term never enters, so a nonzero coefficient + would be silently ignored rather than applied. ``--use-kl-loss`` with + ``--kl-loss-coef`` remains available: that one is a separate loss term, not + a reward modification. + """ + + requires_global_token_loss: bool = False + """Set when the objective is only correct under + ``--calculate-per-token-loss``. + + The default per-sample token-mean reducer divides each sample by its own + response length, which reweights unequal-length responses by + ``1 / response_length``. An objective with no ratio correction has nothing + to absorb that reweighting, so it must be normalised by the global number + of valid response tokens instead. + """ + + requires_on_policy_updates: bool = False + """Set when every optimizer step must consume exactly the rollout that + produced it. + + An objective without an importance-ratio correction cannot account for the + policy having moved, so *five* separate configuration knobs have to agree: + no ``--fully-async`` / ``--hybrid``, ``--max-staleness 0``, + ``--num-steps-per-rollout 1``, ``rollout_batch_size * + n_samples_per_prompt == global_batch_size``, and neither + ``--partial-rollout`` nor ``--use-dynamic-global-batch-size`` (both let the + effective batch size drift at runtime). + + They are one field rather than five because they have one cause. RLOO is + currently the only member, so the bundling is a guess about the next + unclipped estimator; if one ever needs four of the five, split this field + then rather than adding an exception to it. + + Note this is a *stronger* statement than "cannot run fully-async": an + algorithm can be sync-only for unrelated reasons (PPO's critic topology, + or an advantage that needs batch-level statistics the async deployment + only sees a slice of). Those get their own field; do not fold them here. + """ + + @property + def is_group_normalized(self) -> bool: + """Whether rewards get normalised per prompt group on the rollout + side.""" + return self.reward_normalizer != "none" + + +# NOTE(dev): explicit dict literal, deliberately not decorator-based registration. +# The advantage formula and the policy loss execute in two different processes +# whose import graphs differ; decorator registration depends on "was this module +# imported?" and silently loses an algorithm when one side misses the import. +ALGORITHM_SPECS: dict[str, AlgorithmSpec] = { + "grpo": AlgorithmSpec( + name="grpo", + reward_normalizer="group_mean_std", + advantage_fn="grpo_broadcast", + policy_loss_fn="ppo_clip", + ), + "gspo": AlgorithmSpec( + name="gspo", + reward_normalizer="group_mean_std", + advantage_fn="grpo_broadcast", + policy_loss_fn="ppo_clip", + kl_level="sequence", + needs_full_log_probs=True, + ), + "sapo": AlgorithmSpec( + name="sapo", + reward_normalizer="group_mean_std", + advantage_fn="grpo_broadcast", + policy_loss_fn="sapo", + ), + "cispo": AlgorithmSpec( + name="cispo", + reward_normalizer="group_mean_std", + advantage_fn="grpo_broadcast", + policy_loss_fn="cispo", + ), + "rloo": AlgorithmSpec( + name="rloo", + # REINFORCE leave-one-out (arXiv:2402.14740). The baseline is the mean + # of the *other* completions in the prompt group, so the reward stage + # differs from GRPO's while the advantage stage -- broadcast the scalar + # over the response tokens -- is the same one. + reward_normalizer="group_leave_one_out", + advantage_fn="grpo_broadcast", + policy_loss_fn="rloo", + requires_rewards_normalization=True, + forbids_normalize_advantages=True, + min_group_size=2, + forbids_reward_side_kl=True, + requires_global_token_loss=True, + requires_on_policy_updates=True, + ), + "ppo": AlgorithmSpec( + name="ppo", + reward_normalizer="none", + advantage_fn="gae", + policy_loss_fn="ppo_clip", + needs_critic=True, + ), + "reinforce_plus_plus": AlgorithmSpec( + name="reinforce_plus_plus", + advantage_normalization="token_global", + reward_normalizer="none", + advantage_fn="reinforce_plus_plus", + policy_loss_fn="ppo_clip", + requires_normalize_advantages=True, + ), + "reinforce_plus_plus_baseline": AlgorithmSpec( + name="reinforce_plus_plus_baseline", + advantage_normalization="token_global", + reward_normalizer="group_mean", + advantage_fn="reinforce_plus_plus_baseline", + policy_loss_fn="ppo_clip", + requires_normalize_advantages=True, + # `_validate_reinforce_plus_plus_args` in `relax/utils/arguments.py` + # already enforces these by hand, and it is a frozen Task 29 contract + # that owns the wording its tests match on. Declaring them here is + # still not redundant: an undeclared field is not neutral, it asserts + # the default -- leaving them out would have this spec state that the + # estimator runs fine with `--disable-rewards-normalization` and a + # group of one, both false. The frozen function runs first on both + # validation paths, so its messages still win. + requires_rewards_normalization=True, + min_group_size=2, + forbids_reward_side_kl=True, + ), +} + + +def get_algorithm(name: str) -> AlgorithmSpec: + """Look up an algorithm spec by its ``--advantage-estimator`` value.""" + try: + return ALGORITHM_SPECS[name] + except KeyError: + available = ", ".join(ALGORITHM_SPECS) + raise KeyError(f"Unknown advantage estimator {name!r}. Available: {available}") from None + + +def list_algorithm_names() -> list[str]: + """All registered algorithm names, in definition order.""" + return list(ALGORITHM_SPECS) diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 170265a24..bce3e4f98 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -9,6 +9,9 @@ from megatron.core import mpu from torch.utils.checkpoint import checkpoint +from relax.algorithms import get_algorithm +from relax.algorithms.advantages import compute_advantages_and_returns as compute_advantages_and_returns_impl +from relax.algorithms.policy import compute_policy_loss_for from relax.utils.distributed_utils import distributed_masked_normalize, distributed_masked_whiten from relax.utils.misc import load_function from relax.utils.opd.opd_utils import ( @@ -21,17 +24,9 @@ from relax.utils.training.ppo_utils import ( calculate_log_probs_and_entropy, compute_approx_kl, - compute_cispo_loss, compute_gspo_kl, compute_log_probs, compute_opsm_mask, - compute_policy_loss, - compute_rloo_loss, - compute_sapo_loss, - get_advantages_and_returns_batch, - get_grpo_returns, - get_reinforce_plus_plus_baseline_advantages, - get_reinforce_plus_plus_returns, ) from relax.utils.types import RolloutBatch @@ -525,10 +520,12 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) `args.advantage_estimator`. This function extracts rewards, log-probs, values, and masks from - `rollout_data`, computes KL divergences, then applies the chosen advantage - estimator. Supported methods: "grpo", "gspo", "sapo", "cispo", "ppo", "reinforce_plus_plus", - and "reinforce_plus_plus_baseline". When `args.normalize_advantages` is - True, advantages are whitened across the data-parallel group using masked + `rollout_data`, computes KL divergences, then dispatches to the estimator + named by `relax.algorithms.spec.ALGORITHM_SPECS[...].advantage_fn`. The + supported methods are whatever that registry holds -- deliberately not + listed here, because keeping algorithm names in prose is the duplication + the registry exists to remove. When `args.normalize_advantages` is True, + advantages are whitened across the data-parallel group using masked statistics. Early returns if both `log_probs` and `values` are None (intermediate @@ -576,56 +573,21 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) for i in range(len(log_probs)) ] - if args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo", "rloo"]: - rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) - returns = get_grpo_returns(rewards, kl) - # TODO: is the copy necessary? - advantages = [r for r in returns] # noqa: C416 - - elif args.advantage_estimator == "ppo": - old_rewards = rewards - rewards = [] - kl_coef = -args.kl_coef - cp_rank = mpu.get_context_parallel_rank() - for reward, k in zip(old_rewards, kl, strict=False): - k *= kl_coef - if cp_rank == 0: - k[-1] += reward - rewards.append(k) - advantages, returns = get_advantages_and_returns_batch( - total_lengths, - response_lengths, - values, - rewards, - args.gamma, - args.lambd, - padded_total_lengths=padded_total_lengths, - ) - - elif args.advantage_estimator == "reinforce_plus_plus": - rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) - returns = get_reinforce_plus_plus_returns( - rewards=rewards, - kl=kl, - loss_masks=loss_masks, - response_lengths=response_lengths, - total_lengths=total_lengths, - kl_coef=args.kl_coef, - gamma=args.gamma, - ) - advantages = [r for r in returns] # noqa: C416 - - elif args.advantage_estimator == "reinforce_plus_plus_baseline": - rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) - advantages = get_reinforce_plus_plus_baseline_advantages( - rewards=rewards, - kl=kl, - loss_masks=loss_masks, - ) - returns = advantages - - else: - raise NotImplementedError(f"advantage_estimator {args.advantage_estimator} is not supported. ") + advantages, returns = compute_advantages_and_returns_impl( + args, + rewards=rewards, + kl=kl, + loss_masks=loss_masks, + response_lengths=response_lengths, + total_lengths=total_lengths, + values=values, + # Only this path can compute it: `maybe_padded_total_lengths` reads + # args.qkv_format plus the VL / unsplit-forward flags, which the + # Advantages deployment does not have. GAE needs it to slice CP shards + # at the padded offsets; omitting it does not raise, it reads the wrong + # token positions. + padded_total_lengths=padded_total_lengths, + ) # Optional pure OPD mode: remove all non-OPD reward contribution. # This keeps only the OPD KL term injected below. @@ -688,14 +650,17 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) assert all_advs.size() == all_masks.size(), ( f"Shape mismatch before whitening: advantages {all_advs.size()}, masks {all_masks.size()}" ) - is_reinforce_plus_plus = args.advantage_estimator in { - "reinforce_plus_plus", - "reinforce_plus_plus_baseline", - } - if is_reinforce_plus_plus or all_masks.numel() > 0: + # Which normalisation to apply is declared by the algorithm registry. + # This used to be a set of estimator names maintained here and a second + # identical set in `policy_loss_function` below, which is exactly the + # kind of pair that drifts: the two have to stay in step because + # token-global normalisation is only correct together with the + # mask-safe reducer. + is_token_global = get_algorithm(args.advantage_estimator).advantage_normalization == "token_global" + if is_token_global or all_masks.numel() > 0: dp_group = mpu.get_data_parallel_group() - if is_reinforce_plus_plus: + if is_token_global: whitened_advs_flat, raw_mean, raw_variance, valid_count = distributed_masked_normalize( all_advs, all_masks, @@ -848,12 +813,12 @@ def policy_loss_function( else: advantages = batch["advantages"] - is_reinforce_plus_plus = args.advantage_estimator in { - "reinforce_plus_plus", - "reinforce_plus_plus_baseline", - } + # Same registry field as `compute_advantages_and_returns` reads: the + # mask-safe reducer is the other half of token-global normalisation, not an + # independent choice. + is_token_global = get_algorithm(args.advantage_estimator).advantage_normalization == "token_global" - if is_reinforce_plus_plus: + if is_token_global: sum_of_sample_mean = _get_reinforce_plus_plus_mask_safe_reducer(sum_of_sample_mean, batch["loss_masks"]) true_on_policy = getattr(args, "true_on_policy_mode", False) @@ -891,7 +856,8 @@ def policy_loss_function( old_log_probs = [lp.detach() for lp in log_probs] # Pre-gather log probs if needed by OPSM or GSPO to avoid duplicate gathering - need_full_log_probs = args.use_opsm or args.advantage_estimator == "gspo" + algorithm = get_algorithm(args.advantage_estimator) + need_full_log_probs = args.use_opsm or algorithm.needs_full_log_probs full_log_probs = None full_old_log_probs = None @@ -950,7 +916,7 @@ def policy_loss_function( ) # Compute KL divergence (GSPO uses sequence-level KL, others use per-token KL) - if args.advantage_estimator == "gspo": + if algorithm.kl_level == "sequence": ppo_kl = compute_gspo_kl( full_log_probs=full_log_probs, full_old_log_probs=full_old_log_probs, @@ -965,27 +931,7 @@ def policy_loss_function( log_probs = torch.cat(log_probs, dim=0) ppo_kl = old_log_probs - log_probs - if args.advantage_estimator == "sapo": - tau_pos = getattr(args, "sapo_tau_pos", 1.0) - tau_neg = getattr(args, "sapo_tau_neg", 1.05) - pg_loss, pg_clipfrac = compute_sapo_loss( - ppo_kl=ppo_kl, advantages=advantages, tau_pos=tau_pos, tau_neg=tau_neg - ) - elif args.advantage_estimator == "cispo": - pg_loss, pg_clipfrac = compute_cispo_loss( - log_probs=log_probs, - ppo_kl=ppo_kl, - advantages=advantages, - eps_clip=args.eps_clip, - eps_clip_high=args.eps_clip_high, - ) - elif args.advantage_estimator == "rloo": - pg_loss, pg_clipfrac = compute_rloo_loss( - log_probs=log_probs, - advantages=advantages, - ) - else: - pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) + pg_loss, pg_clipfrac = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) if args.use_opsm: pg_loss = pg_loss * opsm_mask @@ -1040,7 +986,7 @@ def policy_loss_function( dynamic_cp_size=batch.get("dynamic_cp_size", None), dynamic_cp_rank=batch.get("dynamic_cp_rank", None), ) - if is_reinforce_plus_plus: + if is_token_global: sum_of_sample_mean = _get_reinforce_plus_plus_mask_safe_reducer( sum_of_sample_mean, modified_response_masks ) @@ -1053,7 +999,7 @@ def policy_loss_function( pg_loss_reducer = custom_pg_loss_reducer_func( total_lengths, response_lengths, pg_loss_masks, args.calculate_per_token_loss ) - if is_reinforce_plus_plus: + if is_token_global: pg_loss_reducer = _get_reinforce_plus_plus_mask_safe_reducer(pg_loss_reducer, pg_loss_masks) else: pg_loss_reducer = sum_of_sample_mean diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 3fd183fc4..4cfb69b75 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -7,23 +7,17 @@ import torch import transfer_queue as tq -from megatron.core import mpu from ray import serve from tensordict import TensorDict +from relax.algorithms.advantages import compute_advantages_and_returns from relax.components.base import Base from relax.utils.async_utils import run as run_ from relax.utils.opd.opd_utils import ( apply_opd_to_advantages, consume_opd_advantage_data, ) -from relax.utils.training.ppo_utils import ( - compute_approx_kl, - get_advantages_and_returns_batch, - get_grpo_returns, - get_reinforce_plus_plus_baseline_advantages, - get_reinforce_plus_plus_returns, -) +from relax.utils.training.ppo_utils import compute_approx_kl @serve.deployment @@ -124,10 +118,9 @@ def compute_advantages_and_returns(self, rollout_data: Dict[str, Any]) -> Dict[s `self.config.advantage_estimator`. This function extracts rewards, log-probs, values, and masks from - `rollout_data`, computes KL divergences, then applies the chosen advantage - estimator. Supported methods: "grpo", "gspo", "sapo", "cispo", - "rloo", "ppo", "reinforce_plus_plus", and - "reinforce_plus_plus_baseline". + `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`). Early returns if both `log_probs` and `values` are None (intermediate pipeline stages). @@ -173,49 +166,19 @@ def compute_advantages_and_returns(self, rollout_data: Dict[str, Any]) -> Dict[s for i in range(len(log_probs)) ] - if self.config.advantage_estimator in ["grpo", "gspo", "sapo", "cispo", "rloo"]: - rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) - returns = get_grpo_returns(rewards, kl) - advantages = list(returns) # make a copy - - elif self.config.advantage_estimator == "ppo": - # TODO: optimize this - old_rewards = rewards - rewards = [] - for reward, k in zip(old_rewards, kl, strict=False): - k *= -self.config.kl_coef - cp_rank = mpu.get_context_parallel_rank() - if cp_rank == 0: - k[-1] += reward - rewards.append(k) - advantages, returns = get_advantages_and_returns_batch( - total_lengths, response_lengths, values, rewards, self.config.gamma, self.config.lambd - ) - - elif self.config.advantage_estimator == "reinforce_plus_plus": - rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) - returns = get_reinforce_plus_plus_returns( - rewards=rewards, - kl=kl, - loss_masks=loss_masks, - response_lengths=response_lengths, - total_lengths=total_lengths, - kl_coef=self.config.kl_coef, - gamma=self.config.gamma, - ) - advantages = list(returns) - - elif self.config.advantage_estimator == "reinforce_plus_plus_baseline": - rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) - advantages = get_reinforce_plus_plus_baseline_advantages( - rewards=rewards, - kl=kl, - loss_masks=loss_masks, - ) - returns = advantages - - else: - raise NotImplementedError(f"advantage_estimator {self.config.advantage_estimator} is not supported. ") + advantages, returns = compute_advantages_and_returns( + self.config, + rewards=rewards, + kl=kl, + loss_masks=loss_masks, + response_lengths=response_lengths, + total_lengths=total_lengths, + values=values, + # `padded_total_lengths` is intentionally omitted: it is derived from + # `args.qkv_format` and the VL / unsplit-forward flags in the Megatron + # worker, and this deployment has no equivalent. Passing nothing keeps + # the behaviour this path already had. + ) # Optional pure OPD mode: remove all non-OPD reward contribution. if getattr(self.config, "use_opd", False) and getattr(self.config, "opd_only_reward", False): diff --git a/relax/core/registry.py b/relax/core/registry.py index a295e3bcb..ce59c24d8 100644 --- a/relax/core/registry.py +++ b/relax/core/registry.py @@ -11,6 +11,7 @@ def __str__(self) -> str: return self.value +from relax.algorithms import ALGORITHM_SPECS from relax.components.actor import Actor from relax.components.actor_fwd import ActorFwd from relax.components.advantages import Advantages @@ -80,68 +81,40 @@ class ROLES_PPO_FULLY_ASYNC_ON_POLICY(StrEnum): reference: str = "reference" -ALGOS = { - "grpo": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "gspo": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "sapo": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "cispo": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "reinforce_plus_plus": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "reinforce_plus_plus_baseline": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "rloo": { - ROLES.rollout: Rollout, - ROLES.actor: Actor, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, - "sft": { - ROLES.sft: SFT, - ROLES.actor: Actor, - }, - "ppo": { +def _rl_roles(*, needs_critic: bool) -> dict: + """Component classes one RL algorithm binds to each role. + + Every RL algorithm starts the same services except for the critic, which + only value-based estimators need. Deriving the table from the registry is + what keeps ``ALGOS`` and ``--advantage-estimator`` from drifting apart: an + estimator argument parsing accepts can no longer fail here at + service-registration time, because both sides read the same dict. + + This decides which class a role maps to, *not* which roles the controller + walks -- that is ``process_role``'s job and it is deliberately left + untouched. ``controller.py`` iterates ``list(process_role(config))`` and + skips any role missing from this dict, so an algorithm without a critic + simply never matches the ``critic`` member the role sets already carry. + """ + roles = { ROLES.rollout: Rollout, ROLES.actor: Actor, - ROLES.critic: Critic, - ROLES.advantages: Advantages, - ROLES.reference: ActorFwd, - ROLES.actor_fwd: ActorFwd, - }, + } + if needs_critic: + roles[ROLES.critic] = Critic + roles[ROLES.advantages] = Advantages + roles[ROLES.reference] = ActorFwd + roles[ROLES.actor_fwd] = ActorFwd + return roles + + +# NOTE(dev): `ALGOS` keys live in a different namespace from AlgorithmSpec names. +# "sft" is selected by `loss_type`, not by `--advantage-estimator`, so it stays a +# separate literal entry rather than being folded into the algorithm registry. +ALGOS = {name: _rl_roles(needs_critic=spec.needs_critic) for name, spec in ALGORITHM_SPECS.items()} +ALGOS["sft"] = { + ROLES.sft: SFT, + ROLES.actor: Actor, } diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index f7d1e710e..64f9dc33d 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -10,6 +10,7 @@ import yaml from sglang_router.launch_router import RouterArgs +from relax.algorithms import get_algorithm, list_algorithm_names from relax.backends.sglang.arguments import sglang_parse_args from relax.backends.sglang.arguments import validate_args as sglang_validate_args from relax.utils import device as device_utils @@ -1727,19 +1728,11 @@ def add_algo_arguments(parser): parser.add_argument( "--advantage-estimator", type=str, - choices=[ - "grpo", - "gspo", - "reinforce_plus_plus", - "reinforce_plus_plus_baseline", - "ppo", - "sapo", - "cispo", - "rloo", - ], + choices=list_algorithm_names(), default="grpo", help=( - "Advantage estimator to use. Note: on-policy distillation (OPD) is now orthogonal " + "Advantage estimator to use. The choices come from the algorithm registry in " + "relax/algorithms/spec.py. Note: on-policy distillation (OPD) is orthogonal " "to the advantage estimator. Use --opd-kl-coef > 0 to enable OPD on top of any estimator. " "'rloo' uses a leave-one-out baseline with an unclipped REINFORCE loss (sync only)." ), @@ -2963,6 +2956,263 @@ def _normalize_sync_ppo_kl_args(args) -> bool: return True +def _assert_spec_implementations_resolve(spec) -> None: + """Check that every implementation the spec names is actually registered. + + Imports the implementation tables lazily: they pull in torch, and this + module is imported for `--help`. + """ + from relax.algorithms.advantages import ADVANTAGE_FNS + from relax.algorithms.policy import POLICY_LOSS_FNS + from relax.algorithms.rewards import REWARD_NORMALIZERS + + for field, key, table in ( + ("reward_normalizer", spec.reward_normalizer, REWARD_NORMALIZERS), + ("advantage_fn", spec.advantage_fn, ADVANTAGE_FNS), + ("policy_loss_fn", spec.policy_loss_fn, POLICY_LOSS_FNS), + ): + if key not in table: + raise ValueError( + f"Algorithm {spec.name!r} declares {field}={key!r}, which is not registered. " + f"Available: {sorted(table)}." + ) + + +def validate_algorithm_args(args) -> None: + """Apply the constraints the algorithm registry declares for this run. + + These rules used to be `if args.advantage_estimator == "..."` checks + scattered across this file, which meant a new algorithm could silently miss + one. They now come from AlgorithmSpec fields, so declaring the algorithm is + enough. Also sets ``args.use_critic``, the only role switch derived from + the algorithm. + + Runs *after* ``_validate_reinforce_plus_plus_args`` on purpose: that + function owns the frozen Task 29 wording for the REINFORCE++ variants, and + checking the same conditions here first would replace its messages. + """ + spec = get_algorithm(args.advantage_estimator) + + # The spec references its implementations by name, so a typo in the registry + # would otherwise surface as a KeyError deep inside a worker on the first + # batch. Resolve them here, while the error can still name the culprit. + _assert_spec_implementations_resolve(spec) + + args.use_critic = spec.needs_critic + + if spec.requires_normalize_advantages and not args.normalize_advantages: + raise ValueError( + f"The {spec.name!r} advantage estimator requires advantage normalization. " + "Please add `--normalize-advantages` to your command." + ) + + if args.n_samples_per_prompt < spec.min_group_size: + raise ValueError( + f"--advantage-estimator {spec.name} requires --n-samples-per-prompt >= {spec.min_group_size} " + f"(got {args.n_samples_per_prompt}); its reward stage is undefined for a smaller group." + ) + + if spec.requires_rewards_normalization and not args.rewards_normalization: + raise ValueError( + f"--advantage-estimator {spec.name} requires rewards normalization to be enabled " + "(its group-reward stage is what --disable-rewards-normalization skips). " + "Please remove --disable-rewards-normalization." + ) + + if spec.forbids_normalize_advantages and args.normalize_advantages: + raise ValueError( + f"--advantage-estimator {spec.name} is incompatible with --normalize-advantages: " + "the latter re-whitens advantages after DP sharding " + "(distributed_masked_whiten in loss.py), which re-introduces the std " + f"normalization {spec.name} removes and makes the result depend on the DP partition. " + "Please remove --normalize-advantages." + ) + + if spec.requires_global_token_loss and not args.calculate_per_token_loss: + raise ValueError( + f"--advantage-estimator {spec.name} requires --calculate-per-token-loss so policy loss is " + "normalized by the global number of valid response tokens. The per-sample token-mean " + "reducer would reweight unequal-length responses by 1 / response_length." + ) + + if spec.requires_on_policy_updates: + if args.fully_async or getattr(args, "hybrid", False): + raise ValueError( + f"--advantage-estimator {spec.name} only supports synchronous (colocate) training. " + "Please remove --fully-async / --hybrid." + ) + if args.max_staleness != 0: + raise ValueError( + f"--advantage-estimator {spec.name} requires --max-staleness 0: the unclipped objective " + "has no importance-ratio correction for stale rollout data." + ) + if args.partial_rollout or args.use_dynamic_global_batch_size: + raise ValueError( + f"--advantage-estimator {spec.name} is incompatible with --partial-rollout / " + "--use-dynamic-global-batch-size: they cause the effective batch size to drift " + "at runtime, breaking the one-update-per-rollout guarantee." + ) + + +def validate_reward_side_kl(args, is_sft: bool) -> None: + """Reject ``--kl-coef`` for estimators that have nowhere to put it. + + Separate from :func:`validate_algorithm_args`, and called much earlier, + because of what runs in between: a nonzero ``--kl-coef`` makes validation + require ``--ref-load`` to exist on disk. Checking this later would report a + missing reference checkpoint for a run whose real problem is that the + estimator would have ignored the coefficient anyway. + """ + if is_sft: + return + spec = get_algorithm(args.advantage_estimator) + if not spec.forbids_reward_side_kl or args.kl_coef == 0: + return + + # `_validate_reinforce_plus_plus_args` is the frozen Task 29 contract and + # owns the wording for its two estimators, but it runs later than this + # point. Give it the first word here rather than pre-empting it -- and only + # on a path that is about to raise anyway, so no other error's precedence + # changes. + _validate_reinforce_plus_plus_args(args, is_sft) + + raise ValueError( + f"--advantage-estimator {spec.name} does not support nonzero --kl-coef: reward-side KL " + "shaping is not implemented for the completion-level signal it trains on. Set --kl-coef 0; " + "for a supported direct KL penalty, provide --ref-load and use --use-kl-loss with " + "--kl-loss-coef." + ) + + +def validate_update_schedule(args) -> None: + """Reject repeated optimizer updates on one rollout. + + Called before ``--num-steps-per-rollout`` is folded into + ``global_batch_size``: afterwards the two are consistent by construction + and a mismatch surfaces as an assertion about batch arithmetic rather than + as the reason the schedule is wrong. + """ + spec = get_algorithm(args.advantage_estimator) + if spec.requires_on_policy_updates and args.num_steps_per_rollout not in (None, 1): + raise ValueError( + f"--advantage-estimator {spec.name} requires --num-steps-per-rollout 1 " + "(the unclipped objective has no ratio correction, so repeated updates on the same " + "rollout would be off-policy)." + ) + + +def derive_global_batch_size(args, *, enforce_consistency: bool = True) -> None: + """Fold ``--num-steps-per-rollout`` into ``global_batch_size``. + + A function rather than three inline lines because + :func:`validate_batch_shape` reads the value it writes, and both have to + run again after ``--custom-config-path`` merges. Leaving the derivation + inline is what made re-running the validator alone *reject a legitimate + config*: a YAML file that switches from ``num_steps_per_rollout: 4`` to + ``1`` should get a global batch of ``rollout * n``, but the validator saw + the stale value derived from 4 and refused it. + """ + if getattr(args, "num_steps_per_rollout", None) is None: + return + global_batch_size = args.rollout_batch_size * args.n_samples_per_prompt // args.num_steps_per_rollout + if enforce_consistency and args.global_batch_size is not None: + assert args.global_batch_size == global_batch_size, ( + f"global_batch_size {args.global_batch_size} is not equal to " + f"rollout_batch_size {args.rollout_batch_size} * n_samples_per_prompt {args.n_samples_per_prompt} " + f"// num_steps_per_rollout {args.num_steps_per_rollout}" + ) + args.global_batch_size = global_batch_size + + +def validate_batch_shape(args) -> None: + """Require the rollout to fill exactly one optimizer step. + + Called after ``global_batch_size`` has taken its final value -- checking + earlier would compare against a number validation is still deriving. + """ + spec = get_algorithm(args.advantage_estimator) + if spec.requires_on_policy_updates and args.rollout_batch_size * args.n_samples_per_prompt != ( + args.global_batch_size + ): + raise ValueError( + f"--advantage-estimator {spec.name} requires exactly one optimizer update per rollout " + "(the unclipped objective has no ratio correction, so a second update on the " + "same rollout is off-policy without correction). This means " + "rollout_batch_size * n_samples_per_prompt must equal global_batch_size, " + f"got {args.rollout_batch_size} * {args.n_samples_per_prompt} = " + f"{args.rollout_batch_size * args.n_samples_per_prompt} != " + f"{args.global_batch_size}." + ) + + +def apply_custom_config_overrides(args) -> None: + """Merge ``--custom-config-path`` YAML into ``args`` and re-check the + result. + + The merge happens late in validation so that a YAML file can override + derived values, which means every algorithm check that already ran was made + against a config we may no longer be training with. Re-running them here is + what stops a YAML file from quietly switching on a flag the algorithm + forbids. + """ + if not args.custom_config_path: + return + + use_critic_before_override = getattr(args, "use_critic", False) + with open(args.custom_config_path) as f: + data = yaml.safe_load(f) or {} + for k, v in data.items(): + if hasattr(args, k): + logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") + setattr(args, k, v) + + if args.loss_type in ("sft", "sft_loss", "sft-loss"): + return + + # Every *algorithm* validator, and the one derivation they read. Be precise + # about the scope: this function does not close every hole, it closes the + # algorithm-shaped ones. + # + # What it covers: `validate_reward_side_kl`, `validate_update_schedule` and + # `validate_batch_shape` were split out of `validate_algorithm_args` + # because argument validation has a derivation order. Re-running only the + # spec-driven validator would leave a YAML file free to 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. `_validate_reinforce_plus_plus_args` is re-run too, because + # the frozen function owns a constraint the spec deliberately does not + # restate. + # + # What it does NOT cover, and a YAML file can still move: the `--ref-load` + # existence check, the `kl_coef`/`kl_loss_coef` exclusion assert, + # `_normalize_sync_ppo_kl_args`, the fully-async resource checks, the + # `rollout_batch_size` derivation, and the over-sampling assert. All of + # them run before the merge and none is an algorithm validator. Closing + # that class properly means merging the YAML *before* validation rather + # than bolting re-runs on after it, which is a larger change than this one. + _validate_reinforce_plus_plus_args(args, is_sft=False) + validate_algorithm_args(args) + validate_reward_side_kl(args, is_sft=False) + validate_update_schedule(args) + # The derivation, then the validator that reads what it writes. Re-running + # the validator alone rejected a legitimate config: a YAML switching + # `num_steps_per_rollout` from 4 to 1 should get a global batch of + # `rollout * n`, and the validator instead saw the value derived from 4. + # `enforce_consistency=False` because the stale value is, by construction, + # the one derived before the merge -- comparing against it is the bug. + derive_global_batch_size(args, enforce_consistency=False) + validate_batch_shape(args) + if args.use_critic != use_critic_before_override: + # Role composition and the offload flags were derived from the pre-override + # value earlier in validation, so accepting the new one here would leave the + # run half-configured rather than either fully critic or fully critic-free. + raise ValueError( + f"--custom-config-path changed the algorithm to {args.advantage_estimator!r}, which needs a different " + f"critic setup than the one already derived. Pass --advantage-estimator on the command line instead " + f"of overriding it from YAML." + ) + + def slime_validate_args(args): # Backward compatibility: old scripts may pass --enable-gloo-process-groups if not hasattr(args, "use_gloo_process_groups"): @@ -3057,13 +3307,7 @@ def slime_validate_args(args): "whereas 'partial_rollout' introduces partial off-policy behavior. These two features are mutually exclusive." ) - if not is_sft and args.advantage_estimator == "rloo" and args.kl_coef != 0: - raise ValueError( - "--advantage-estimator rloo does not support nonzero --kl-coef: reward-side KL shaping " - "is not implemented for the completion-level leave-one-out signal. Set --kl-coef 0; " - "for a supported direct KL penalty, provide --ref-load and use --use-kl-loss with " - "--kl-loss-coef." - ) + validate_reward_side_kl(args, is_sft) if not is_sft and (args.kl_coef != 0 or args.use_kl_loss): if not os.path.exists(args.ref_load): @@ -3148,10 +3392,14 @@ def slime_validate_args(args): raise ValueError("Either --rollout-batch-size or --global-batch-size must be set.") if args.n_samples_per_prompt <= 0: raise ValueError("--n-samples-per-prompt must be positive when deriving --rollout-batch-size.") - if args.advantage_estimator == "rloo" and args.global_batch_size % args.n_samples_per_prompt != 0: + # An estimator that must consume exactly one rollout per update cannot + # absorb the remainder this floor division would drop. + if get_algorithm(args.advantage_estimator).requires_on_policy_updates and ( + args.global_batch_size % args.n_samples_per_prompt != 0 + ): raise ValueError( - "--global-batch-size must be divisible by --n-samples-per-prompt for RLOO when " - "--rollout-batch-size is omitted, got " + f"--global-batch-size must be divisible by --n-samples-per-prompt for " + f"{args.advantage_estimator} when --rollout-batch-size is omitted, got " f"{args.global_batch_size} % {args.n_samples_per_prompt} != 0." ) args.rollout_batch_size = args.global_batch_size // args.n_samples_per_prompt @@ -3161,11 +3409,8 @@ def slime_validate_args(args): ) if not is_sft: - if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]: - assert args.normalize_advantages, ( - "The 'reinforce_plus_plus' and 'reinforce_plus_plus_baseline' advantage estimators " - "require advantage normalization. Please add `--normalize-advantages` to your command." - ) + validate_algorithm_args(args) + if args.fully_async: assert not args.normalize_advantages, ( "Advantage normalization is not supported in fully-async mode (--fully-async). " @@ -3316,7 +3561,9 @@ def slime_validate_args(args): logger.info("--loss-type sft: auto-enabling --balance-data for DP-balanced batching.") args.balance_data = True - args.use_critic = args.advantage_estimator == "ppo" + # `use_critic` is set by validate_algorithm_args for RL runs; SFT never has one. + if is_sft: + args.use_critic = False # Synchronous PPO has no producer for # `ref_log_probs`: actor's ref forward in backends/megatron/actor.py:800 is # gated on `advantage_estimator != "ppo"`, and the sync role set does not @@ -3483,75 +3730,13 @@ def slime_validate_args(args): if args.eval_function_path is None: args.eval_function_path = args.rollout_function_path - if args.advantage_estimator == "rloo" and args.num_steps_per_rollout not in (None, 1): - raise ValueError( - "--advantage-estimator rloo requires --num-steps-per-rollout 1 " - "(the unclipped objective has no ratio correction, so repeated updates on the same " - "rollout would be off-policy)." - ) + if not is_sft: + validate_update_schedule(args) - if args.num_steps_per_rollout is not None: - global_batch_size = args.rollout_batch_size * args.n_samples_per_prompt // args.num_steps_per_rollout - if args.global_batch_size is not None: - assert args.global_batch_size == global_batch_size, ( - f"global_batch_size {args.global_batch_size} is not equal to " - f"rollout_batch_size {args.rollout_batch_size} * n_samples_per_prompt {args.n_samples_per_prompt} " - f"// num_steps_per_rollout {args.num_steps_per_rollout}" - ) - args.global_batch_size = global_batch_size + derive_global_batch_size(args) - if args.advantage_estimator == "rloo": - if args.n_samples_per_prompt < 2: - raise ValueError( - "--advantage-estimator rloo requires --n-samples-per-prompt >= 2 " - "(the leave-one-out baseline divides by G-1; G=1 is undefined)." - ) - if args.fully_async or getattr(args, "hybrid", False): - raise ValueError( - "--advantage-estimator rloo only supports synchronous (colocate) training. " - "Please remove --fully-async / --hybrid." - ) - if not args.calculate_per_token_loss: - raise ValueError( - "--advantage-estimator rloo requires --calculate-per-token-loss so policy loss is " - "normalized by the global number of valid response tokens. The per-sample token-mean " - "reducer would reweight unequal-length responses by 1 / response_length." - ) - if args.max_staleness != 0: - raise ValueError( - "--advantage-estimator rloo requires --max-staleness 0: the unclipped objective has no " - "importance-ratio correction for stale rollout data." - ) - if not args.rewards_normalization: - raise ValueError( - "--advantage-estimator rloo requires rewards normalization to be enabled " - "(the leave-one-out baseline is computed inside the group-reward path that " - "--disable-rewards-normalization skips). Please remove --disable-rewards-normalization." - ) - if args.normalize_advantages: - raise ValueError( - "--advantage-estimator rloo is incompatible with --normalize-advantages: " - "the latter re-whitens advantages after DP sharding " - "(distributed_masked_whiten in loss.py), which re-introduces the std " - "normalization RLOO removes and makes the result depend on the DP partition. " - "Please remove --normalize-advantages." - ) - if args.rollout_batch_size * args.n_samples_per_prompt != args.global_batch_size: - raise ValueError( - "--advantage-estimator rloo requires exactly one optimizer update per rollout " - "(the unclipped objective has no ratio correction, so a second update on the " - "same rollout is off-policy without correction). This means " - "rollout_batch_size * n_samples_per_prompt must equal global_batch_size, " - f"got {args.rollout_batch_size} * {args.n_samples_per_prompt} = " - f"{args.rollout_batch_size * args.n_samples_per_prompt} != " - f"{args.global_batch_size}." - ) - if args.partial_rollout or args.use_dynamic_global_batch_size: - raise ValueError( - "--advantage-estimator rloo is incompatible with --partial-rollout / " - "--use-dynamic-global-batch-size: they cause the effective batch size to drift " - "at runtime, breaking the one-update-per-rollout guarantee." - ) + if not is_sft: + validate_batch_shape(args) if args.n_samples_per_prompt == 1: args.grpo_std_normalization = False @@ -3625,13 +3810,7 @@ def slime_validate_args(args): if args.use_rollout_routing_replay: args.use_routing_replay = True - if args.custom_config_path: - with open(args.custom_config_path) as f: - data = yaml.safe_load(f) or {} - for k, v in data.items(): - if hasattr(args, k): - logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") - setattr(args, k, v) + apply_custom_config_overrides(args) if args.eval_max_context_len is None: logger.info( diff --git a/relax/utils/metrics/metric_utils.py b/relax/utils/metrics/metric_utils.py index e0994262f..f5a28e4d0 100644 --- a/relax/utils/metrics/metric_utils.py +++ b/relax/utils/metrics/metric_utils.py @@ -5,6 +5,7 @@ import numpy as np import torch +from relax.algorithms.spec import get_algorithm from relax.utils.types import Sample @@ -111,7 +112,13 @@ def finalize_rollout_explicit_metric_values(metric_values: dict[str, list[float] def _compute_rloo_group_diagnostics(args, samples: list[Sample]) -> dict[str, float]: - """Compute RLOO-specific diagnostics from training rollout samples. + """Compute leave-one-out diagnostics from training rollout samples. + + Gated on the algorithm registry's ``reward_normalizer`` rather than on the + estimator's name: what makes these numbers meaningful is that the reward + stage produced a leave-one-out baseline, so any algorithm declaring that + normalizer gets them. The ``rloo/`` key prefix stays as-is -- those metric + names are already published to dashboards. These keys are returned with an ``rloo/`` prefix. The training rollout logger adds the outer ``rollout/`` prefix before publishing them. @@ -124,7 +131,8 @@ def _compute_rloo_group_diagnostics(args, samples: list[Sample]) -> dict[str, fl length and therefore remains distinct from a sample whose response exists but is fully masked. """ - if getattr(args, "advantage_estimator", None) != "rloo": + spec = get_algorithm(getattr(args, "advantage_estimator", None) or "grpo") + if spec.reward_normalizer != "group_leave_one_out": return {} if ( getattr(args, "custom_reward_post_process_path", None) is not None diff --git a/relax/utils/utils.py b/relax/utils/utils.py index d7f5e93f3..902bc42c4 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -11,11 +11,12 @@ import torch from tensordict import TensorDict +from relax.algorithms import get_algorithm +from relax.algorithms.rewards import REWARD_NORMALIZERS from relax.utils.device import get_ray_accelerator_name from relax.utils.env import Envs, validate_env from relax.utils.logging_utils import get_logger from relax.utils.misc import load_function -from relax.utils.training.ppo_utils import compute_rloo_leave_one_out_rewards from relax.utils.types import Sample @@ -133,9 +134,14 @@ def convert_samples_to_train_data(args: Any, samples: list[Sample] | list[list[S # overwriting the raw reward # populate this field for a subset of samples (e.g. SWE but not code). if any(sample.metadata and "raw_reward" in sample.metadata for sample in samples): + # NOTE(dev): the fallback must stay a scalar. `sample.reward` is a dict + # whenever the reward function returns named components (any run using + # --reward-key), and mixing dicts into this column + # makes dict_to_tensordict raise. `raw_rewards` already holds the scalar + # that post_process_rewards selected for each sample. train_data["raw_reward"] = [ - sample.metadata["raw_reward"] if sample.metadata and "raw_reward" in sample.metadata else sample.reward - for sample in samples + sample.metadata["raw_reward"] if sample.metadata and "raw_reward" in sample.metadata else raw_reward + for sample, raw_reward in zip(samples, raw_rewards, strict=True) ] # For rollout buffer @@ -180,50 +186,21 @@ def post_process_rewards(args: Any, samples: list[Sample] | list[list[Sample]]): return custom_reward_post_process_func(args, samples) raw_rewards = [sample.get_reward_value(args) for sample in samples] + # Second short-circuit: this one replaces the normalizer wholesale. Any + # algorithm whose reward stage is load-bearing would be silently skipped + # here while the run still reports itself as that algorithm; none of the + # currently registered normalizers is, so this stays as it was. if getattr(args, "agentic_custom_advantage_path", None) is not None: return raw_rewards, [sample.custom_advantage for sample in samples] - if ( - args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline", "rloo"] - and args.rewards_normalization - ): - # group norm - rewards = torch.tensor(raw_rewards, dtype=torch.float) - positions_by_group: dict[int, list[int]] = {} - for position, sample in enumerate(samples): - if sample.group_index is None: - raise ValueError("Sample.group_index is required for group reward normalization.") - if sample.group_index not in positions_by_group: - positions_by_group[sample.group_index] = [] - positions_by_group[sample.group_index].append(position) - - normalized_rewards = torch.empty_like(rewards) - for group_index, positions in positions_by_group.items(): - if len(positions) != args.n_samples_per_prompt: - raise ValueError( - f"Reward group {group_index} has {len(positions)} samples, expected {args.n_samples_per_prompt}." - ) - group_rewards = rewards[positions] - if args.advantage_estimator == "rloo": - finite_mask = torch.isfinite(group_rewards) - if not finite_mask.all(): - invalid_group_positions = (~finite_mask).nonzero(as_tuple=False).flatten().tolist() - invalid_sample_positions = [positions[position] for position in invalid_group_positions] - invalid_values = group_rewards[~finite_mask].tolist() - raise ValueError( - f"RLOO group_index={group_index} contains non-finite reward(s) at " - f"group position(s) {invalid_group_positions}, sample position(s) " - f"{invalid_sample_positions}: {invalid_values}." - ) - group_rewards = compute_rloo_leave_one_out_rewards(group_rewards) - else: - group_rewards = group_rewards - group_rewards.mean() - if args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo"] and args.grpo_std_normalization: - group_rewards = group_rewards / (group_rewards.std() + 1e-6) - normalized_rewards[positions] = group_rewards - return raw_rewards, normalized_rewards.tolist() + if not args.rewards_normalization: + return raw_rewards, raw_rewards - return raw_rewards, raw_rewards + # Which normalization to apply is declared by the algorithm registry rather + # than by a whitelist of estimator names maintained here. + spec = get_algorithm(args.advantage_estimator) + normalizer = REWARD_NORMALIZERS[spec.reward_normalizer] + return raw_rewards, normalizer(args, samples, raw_rewards) def dict_to_tensordict( @@ -451,7 +428,7 @@ def get_debug_data(args, rollout_id: int, batch_size, dp_rank: int) -> Dict[str, original_num_rows = len(data) if ( args.custom_reward_post_process_path is None - and args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline", "rloo"] + and get_algorithm(args.advantage_estimator).is_group_normalized and args.rewards_normalization ): group_ids = list(dict.fromkeys(sample.group_index for sample in data)) diff --git a/tests/algorithms/__init__.py b/tests/algorithms/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/algorithms/test_advantage_estimators.py b/tests/algorithms/test_advantage_estimators.py new file mode 100644 index 000000000..08a42f822 --- /dev/null +++ b/tests/algorithms/test_advantage_estimators.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Golden-value tests for the extracted advantage estimators.""" + +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + +from relax.algorithms.advantages import ADVANTAGE_FNS, compute_advantages_and_returns # noqa: E402 + + +def _args(estimator, **overrides): + base = dict(advantage_estimator=estimator, kl_coef=0.0, gamma=1.0, lambd=1.0) + base.update(overrides) + return SimpleNamespace(**base) + + +def _inputs(lengths=(3, 2)): + return dict( + kl=[torch.zeros(n, dtype=torch.float32) for n in lengths], + loss_masks=[torch.ones(n, dtype=torch.float32) for n in lengths], + response_lengths=list(lengths), + total_lengths=[n + 2 for n in lengths], + values=None, + ) + + +# ---------------- dispatch ---------------- + + +def test_registry_covers_every_spec_advantage_id(): + from relax.algorithms import list_algorithm_names + from relax.algorithms.spec import get_algorithm + + for name in list_algorithm_names(): + assert get_algorithm(name).advantage_fn in ADVANTAGE_FNS + + +def test_unknown_advantage_fn_is_not_silently_tolerated(): + with pytest.raises(KeyError): + ADVANTAGE_FNS["not_a_real_fn"] + + +# ---------------- grpo family ---------------- + + +def test_grpo_broadcast_repeats_the_scalar_over_tokens(): + adv, ret = compute_advantages_and_returns(_args("grpo"), rewards=[1.5, -2.0], **_inputs()) + assert torch.equal(adv[0], torch.full((3,), 1.5)) + assert torch.equal(adv[1], torch.full((2,), -2.0)) + assert torch.equal(ret[0], adv[0]) + + +def test_grpo_accepts_a_tensor_as_well_as_a_list(): + from_list = compute_advantages_and_returns(_args("grpo"), rewards=[1.5, -2.0], **_inputs())[0] + from_tensor = compute_advantages_and_returns(_args("grpo"), rewards=torch.tensor([1.5, -2.0]), **_inputs())[0] + assert torch.equal(from_list[0], from_tensor[0]) + + +def test_grpo_advantages_is_a_distinct_list_from_returns(): + """Legacy did `advantages = list(returns)`; rebinding one must not touch the other.""" + adv, ret = compute_advantages_and_returns(_args("grpo"), rewards=[1.0, 1.0], **_inputs()) + assert adv is not ret + adv[0] = torch.zeros(3) + assert not torch.equal(ret[0], adv[0]) + + +@pytest.mark.parametrize("estimator", ["grpo", "gspo", "sapo", "cispo"]) +def test_grpo_family_produces_identical_advantages(estimator): + baseline, _ = compute_advantages_and_returns(_args("grpo"), rewards=[0.5, -0.5], **_inputs()) + actual, _ = compute_advantages_and_returns(_args(estimator), rewards=[0.5, -0.5], **_inputs()) + for left, right in zip(baseline, actual, strict=True): + assert torch.equal(left, right) + + +# ---------------- reinforce++ baseline ---------------- + + +def test_reinforce_plus_plus_baseline_aliases_returns_to_advantages(): + """Legacy did `returns = advantages` (the same list object). Preserve it.""" + adv, ret = compute_advantages_and_returns(_args("reinforce_plus_plus_baseline"), rewards=[1.0, 1.0], **_inputs()) + assert ret is adv + + +def test_reinforce_plus_plus_baseline_keeps_kl_out_of_the_advantage(): + """The baseline variant regularises via a separate k2 loss, not via the + advantage, so the KL tensor only supplies the per-token shape.""" + inputs = _inputs(lengths=(2,)) + inputs["kl"] = [torch.tensor([2.0, 4.0])] + inputs["loss_masks"] = [torch.ones(2)] + adv, _ = compute_advantages_and_returns( + _args("reinforce_plus_plus_baseline", kl_coef=0.5), rewards=[3.0], **inputs + ) + # The scalar reward is broadcast to every unmasked token; kl is not subtracted + # even though kl_coef is non-zero (argument validation forbids that combination). + assert torch.equal(adv[0], torch.tensor([3.0, 3.0])) + + +def test_reinforce_plus_plus_baseline_zeroes_masked_tokens(): + inputs = _inputs(lengths=(3,)) + inputs["kl"] = [torch.zeros(3)] + inputs["loss_masks"] = [torch.tensor([1.0, 0.0, 1.0])] + adv, _ = compute_advantages_and_returns(_args("reinforce_plus_plus_baseline"), rewards=[2.0], **inputs) + assert torch.equal(adv[0], torch.tensor([2.0, 0.0, 2.0])) diff --git a/tests/algorithms/test_algorithm_registry.py b/tests/algorithms/test_algorithm_registry.py new file mode 100644 index 000000000..5e8706385 --- /dev/null +++ b/tests/algorithms/test_algorithm_registry.py @@ -0,0 +1,155 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for the algorithm registry.""" + +import inspect + +import pytest + +from relax.algorithms import get_algorithm, list_algorithm_names +from relax.algorithms.spec import ALGORITHM_SPECS, AlgorithmSpec + + +EXPECTED_NAMES = [ + "grpo", + "gspo", + "sapo", + "cispo", + "ppo", + "reinforce_plus_plus", + "reinforce_plus_plus_baseline", +] + + +def test_all_expected_algorithms_registered(): + for name in EXPECTED_NAMES: + assert name in ALGORITHM_SPECS, f"{name} missing from ALGORITHM_SPECS" + + +def test_spec_name_matches_dict_key(): + for key, spec in ALGORITHM_SPECS.items(): + assert spec.name == key + + +def test_spec_is_frozen(): + spec = get_algorithm("grpo") + with pytest.raises(Exception): + spec.name = "mutated" + + +def test_get_algorithm_unknown_name_raises_with_available_names(): + with pytest.raises(KeyError) as exc: + get_algorithm("does_not_exist") + assert "grpo" in str(exc.value) + + +def test_list_algorithm_names_matches_registry_keys(): + assert list_algorithm_names() == list(ALGORITHM_SPECS.keys()) + + +def test_grpo_family_shares_one_advantage_fn(): + """grpo/gspo/sapo/cispo are identical at the advantage layer.""" + ids = {get_algorithm(n).advantage_fn for n in ("grpo", "gspo", "sapo", "cispo")} + assert ids == {"grpo_broadcast"} + + +def test_reward_normalizer_ids_match_current_behavior(): + for name in ("grpo", "gspo", "sapo", "cispo"): + assert get_algorithm(name).reward_normalizer == "group_mean_std" + assert get_algorithm("reinforce_plus_plus_baseline").reward_normalizer == "group_mean" + for name in ("ppo", "reinforce_plus_plus"): + assert get_algorithm(name).reward_normalizer == "none" + + +def test_is_group_normalized_matches_the_legacy_whitelist(): + """The pre-registry whitelist, exactly. + + group. + """ + legacy = {"grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline"} + actual = {n for n in EXPECTED_NAMES if get_algorithm(n).is_group_normalized} + assert actual == legacy + + +def test_gspo_is_the_only_sequence_level_kl(): + seq = {n for n in EXPECTED_NAMES if get_algorithm(n).kl_level == "sequence"} + assert seq == {"gspo"} + + +def test_gspo_is_the_only_one_needing_full_log_probs(): + need = {n for n in EXPECTED_NAMES if get_algorithm(n).needs_full_log_probs} + assert need == {"gspo"} + + +def test_ppo_is_the_only_algorithm_needing_a_critic(): + """`needs_critic` is what `relax/core/registry.py` reads to decide whether + ALGOS binds the Critic component, so this set is load-bearing.""" + critic = {n for n in EXPECTED_NAMES if get_algorithm(n).needs_critic} + assert critic == {"ppo"} + + +def test_reinforce_family_requires_normalize_advantages(): + for name in ("reinforce_plus_plus", "reinforce_plus_plus_baseline"): + assert get_algorithm(name).requires_normalize_advantages is True + assert get_algorithm("grpo").requires_normalize_advantages is False + + +def test_policy_loss_ids_match_current_behavior(): + assert get_algorithm("sapo").policy_loss_fn == "sapo" + assert get_algorithm("cispo").policy_loss_fn == "cispo" + for name in ("grpo", "gspo", "ppo", "reinforce_plus_plus", "reinforce_plus_plus_baseline"): + assert get_algorithm(name).policy_loss_fn == "ppo_clip" + + +def test_defaults_are_permissive(): + spec = AlgorithmSpec(name="x", reward_normalizer="none", advantage_fn="a", policy_loss_fn="ppo_clip") + assert spec.kl_level == "token" + assert spec.needs_full_log_probs is False + assert spec.needs_critic is False + assert spec.requires_normalize_advantages is False + + +def test_spec_module_has_no_heavy_imports(): + """The registry must import on a CPU-only runner with just torch + available.""" + import relax.algorithms.spec as spec_mod + + src = inspect.getsource(spec_mod) + for banned in ( + "import megatron", + "from megatron", + "import ray", + "from ray", + "import transfer_queue", + "import tensordict", + "from relax.components", + "from relax.backends", + ): + assert banned not in src, f"spec.py must not import {banned}" + + +def test_every_spec_identifier_resolves_to_a_registered_implementation(): + """A typo in the registry must not wait until the first batch to + surface.""" + from relax.algorithms.advantages import ADVANTAGE_FNS + from relax.algorithms.policy import POLICY_LOSS_FNS + from relax.algorithms.rewards import REWARD_NORMALIZERS + + for name in list_algorithm_names(): + spec = get_algorithm(name) + assert spec.reward_normalizer in REWARD_NORMALIZERS, name + assert spec.advantage_fn in ADVANTAGE_FNS, name + assert spec.policy_loss_fn in POLICY_LOSS_FNS, name + + +def test_every_spec_declares_a_kl_level_the_loss_knows_how_to_read(): + """``kl_level`` has no dispatch table to fail against. + + ``relax/backends/megatron/loss.py`` reads it as ``== "sequence"``, so a + misspelled value does not raise -- it silently selects token-level KL and + the run trains the wrong objective while reporting the right algorithm + name. The three fields above cannot fail that way because a bad key raises + on lookup; this one needs the check written out. + """ + for name in list_algorithm_names(): + assert get_algorithm(name).kl_level in ("token", "sequence"), name diff --git a/tests/algorithms/test_algos_roles.py b/tests/algorithms/test_algos_roles.py new file mode 100644 index 000000000..d5320a24b --- /dev/null +++ b/tests/algorithms/test_algos_roles.py @@ -0,0 +1,135 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""ALGOS must cover every registered algorithm, with the right roles. + +`ALGOS` was a hand-written dict, one literal block per algorithm. That is the +kind of table an estimator can be accepted by argparse yet be missing from -- +`reinforce_plus_plus` and `reinforce_plus_plus_baseline` were exactly that for a +while, crashing `controller.register_all_serve` with `ValueError: Algorithm key +'...' not registered in ALGOS` until they were added by hand. Deriving the table +from the registry removes the class of bug rather than one instance of it. + +The roles are not uniform: PPO needs a critic and the policy-gradient +estimators do not, so the derivation reads `AlgorithmSpec.needs_critic` instead +of handing every algorithm the same set. These tests pin both halves. + +`relax.core.registry` imports the component classes, which import megatron, so +the behavioural checks are gated on that. The source-level checks are not: they +are the ones that run on the CPU-only CI runner, and they are what catches a +regression back to hand-written entries. +""" + +import pathlib + +import pytest + + +REGISTRY_PATH = pathlib.Path(__file__).resolve().parents[2] / "relax" / "core" / "registry.py" + +try: + from relax.core.registry import ALGOS, ROLES + + HAS_MEGATRON = True +except ImportError: # pragma: no cover - depends on the runner + ALGOS = ROLES = None + HAS_MEGATRON = False + +requires_megatron = pytest.mark.skipif(not HAS_MEGATRON, reason="relax.core.registry requires megatron") + + +# ---------------- runs everywhere, including CPU-only CI ---------------- + + +def test_algos_is_derived_from_the_registry_not_hand_written(): + src = REGISTRY_PATH.read_text(encoding="utf-8") + + assert "ALGORITHM_SPECS.items()" in src, "ALGOS no longer derives its entries from the registry" + for name in ("grpo", "gspo", "sapo", "cispo", "ppo"): + assert f'"{name}": {{' not in src, f"ALGOS hand-writes a role dict for {name} again" + + +def test_role_topology_is_driven_by_needs_critic(): + """Which component classes an algorithm binds must come from the spec. + + Scoped to the ALGOS derivation on purpose. `process_role` still branches on + the literal "ppo" to pick the role *iteration order*; that is the + controller's orchestration surface, it is untouched by this change, and + folding it in would be a separate proposal. + """ + src = REGISTRY_PATH.read_text(encoding="utf-8") + + assert "needs_critic=spec.needs_critic" in src, "ALGOS no longer derives the critic role from AlgorithmSpec" + + +def test_sft_stays_a_separate_literal_entry(): + """SFT is selected by loss_type, so it must not come from the estimator + registry.""" + src = REGISTRY_PATH.read_text(encoding="utf-8") + assert 'ALGOS["sft"]' in src + + +def test_sft_is_not_an_estimator(): + from relax.algorithms import list_algorithm_names + + assert "sft" not in list_algorithm_names() + + +# ---------------- needs the real component classes ---------------- + + +@requires_megatron +def test_every_registered_algorithm_has_a_role_mapping(): + """The regression that used to crash the controller at startup.""" + from relax.algorithms import list_algorithm_names + + missing = [name for name in list_algorithm_names() if name not in ALGOS] + assert not missing, f"{missing} would raise ValueError in controller.register_all_serve" + + +@requires_megatron +@pytest.mark.parametrize("name", ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]) +def test_previously_missing_algorithms_are_covered(name): + """These two are why the derivation exists.""" + assert name in ALGOS + + +@requires_megatron +def test_role_set_differs_only_by_the_critic(): + from relax.algorithms import ALGORITHM_SPECS + + base = {ROLES.rollout, ROLES.actor, ROLES.advantages, ROLES.reference, ROLES.actor_fwd} + for name, spec in ALGORITHM_SPECS.items(): + expected = base | {ROLES.critic} if spec.needs_critic else base + assert set(ALGOS[name]) == expected, f"{name} has unexpected roles" + + +@requires_megatron +def test_ppo_keeps_its_critic(): + """PPO is the one value-based estimator; dropping Critic here would leave + `controller.register_all_serve` skipping the role and the run training + without a value function.""" + from relax.components.critic import Critic + + assert ALGOS["ppo"][ROLES.critic] is Critic + + +@requires_megatron +def test_policy_gradient_algorithms_have_no_critic(): + for name in ("grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus"): + assert ROLES.critic not in ALGOS[name], f"{name} would start a critic service it never uses" + + +@requires_megatron +def test_sft_roles_are_unchanged(): + assert set(ALGOS["sft"]) == {ROLES.sft, ROLES.actor} + + +@requires_megatron +def test_each_algorithm_gets_an_independent_role_dict(): + """controller.py copies and then mutates these; sharing one dict object + would leak an optional role from one algorithm into all the others.""" + from relax.algorithms import list_algorithm_names + + names = list_algorithm_names() + for left, right in zip(names, names[1:], strict=False): + assert ALGOS[left] is not ALGOS[right], f"{left} and {right} share one dict object" diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py new file mode 100644 index 000000000..7c5eec96d --- /dev/null +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -0,0 +1,316 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Argument parsing and validation must read the registry, not string lists.""" + +import argparse +import importlib +import pathlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + + +ARGS_PATH = pathlib.Path(__file__).resolve().parents[2] / "relax" / "utils" / "arguments.py" + + +@pytest.fixture() +def arguments_module(monkeypatch): + """Import relax.utils.arguments with its heavy optional deps stubbed out. + + Mirrors tests/utils/test_arguments_opd_teacher_colocate.py. + """ + router_pkg = ModuleType("sglang_router") + launch_router = ModuleType("sglang_router.launch_router") + launch_router.RouterArgs = object + monkeypatch.setitem(sys.modules, "sglang_router", router_pkg) + monkeypatch.setitem(sys.modules, "sglang_router.launch_router", launch_router) + + sglang_arguments = ModuleType("relax.backends.sglang.arguments") + sglang_arguments.sglang_parse_args = lambda: None + sglang_arguments.validate_args = lambda args: args + monkeypatch.setitem(sys.modules, "relax.backends.sglang.arguments", sglang_arguments) + + device = ModuleType("relax.utils.device") + device.get_dist_backend = lambda: "gloo" + monkeypatch.setitem(sys.modules, "relax.utils.device", device) + + eval_config = ModuleType("relax.utils.training.eval_config") + eval_config.EvalDatasetConfig = dict + eval_config.build_eval_dataset_configs = lambda args, datasets_config, defaults: [] + eval_config.build_named_prompt_data_configs = lambda values: [] + eval_config.ensure_dataset_list = lambda values: values or [] + monkeypatch.setitem(sys.modules, "relax.utils.training.eval_config", eval_config) + + sys.modules.pop("relax.utils.arguments", None) + module = importlib.import_module("relax.utils.arguments") + yield module + sys.modules.pop("relax.utils.arguments", None) + + +def _args(estimator="grpo", **overrides): + base = dict( + advantage_estimator=estimator, + normalize_advantages=False, + rewards_normalization=True, + custom_reward_post_process_path=None, + n_samples_per_prompt=4, + reward_key="score", + use_critic=False, + fully_async=False, + hybrid=False, + dynamic_sampling_filter_path=None, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +# ---------------- source-level: no hardcoded name lists ---------------- + + +def test_choices_come_from_the_registry(): + assert "choices=list_algorithm_names()" in ARGS_PATH.read_text(encoding="utf-8") + + +def test_no_hardcoded_estimator_choice_list(): + src = ARGS_PATH.read_text(encoding="utf-8") + assert '"reinforce_plus_plus_baseline",\n "ppo",' not in src + + +def test_no_estimator_name_comparisons_remain(): + src = ARGS_PATH.read_text(encoding="utf-8") + for banned in ( + 'args.advantage_estimator == "ppo"', + 'args.advantage_estimator in ["reinforce_plus_plus"', + ): + assert banned not in src, f"arguments.py still contains: {banned}" + + +def test_validation_reads_spec_fields(): + src = ARGS_PATH.read_text(encoding="utf-8") + for field in ( + "needs_critic", + "requires_normalize_advantages", + ): + assert field in src, f"arguments.py does not consult spec.{field}" + + +# ---------------- behaviour ---------------- + + +def test_parser_rejects_an_unregistered_estimator(arguments_module): + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + + with pytest.raises(SystemExit): + parser.parse_args(["--advantage-estimator", "not_an_algorithm"]) + + +def test_every_registered_algorithm_is_an_accepted_choice(arguments_module): + from relax.algorithms import list_algorithm_names + + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + + for name in list_algorithm_names(): + parsed = parser.parse_args(["--advantage-estimator", name]) + assert parsed.advantage_estimator == name + + +def test_reinforce_family_requires_normalize_advantages(arguments_module): + for estimator in ("reinforce_plus_plus", "reinforce_plus_plus_baseline"): + with pytest.raises(ValueError, match="normalize-advantages"): + arguments_module.validate_algorithm_args(_args(estimator, normalize_advantages=False)) + + +def test_reinforce_family_passes_with_normalize_advantages(arguments_module): + args = _args("reinforce_plus_plus", normalize_advantages=True) + arguments_module.validate_algorithm_args(args) + assert args.use_critic is False + + +def test_ppo_is_runnable_and_turns_on_the_critic(arguments_module): + """PPO is enabled upstream again; `use_critic` is the switch that makes + `relax/core/registry.py` bind the Critic component.""" + args = _args("ppo", reward_key=None) + arguments_module.validate_algorithm_args(args) + assert args.use_critic is True + + +@pytest.mark.parametrize("estimator", ["grpo", "gspo", "sapo", "cispo"]) +def test_grpo_family_passes_with_defaults(arguments_module, estimator): + args = _args(estimator, reward_key=None) + arguments_module.validate_algorithm_args(args) + assert args.use_critic is False + + +def test_algorithms_do_not_police_reward_key(arguments_module): + """No currently registered algorithm constrains --reward-key.""" + args = _args("grpo", reward_key=None) + arguments_module.validate_algorithm_args(args) + + +def test_unknown_estimator_raises_from_the_registry(arguments_module): + with pytest.raises(KeyError, match="Unknown advantage estimator"): + arguments_module.validate_algorithm_args(_args("not_an_algorithm")) + + +# ---------------- --custom-config-path override timing ---------------- + + +def _write_yaml(tmp_path, body): + path = tmp_path / "override.yaml" + path.write_text(body, encoding="utf-8") + return str(path) + + +def _overridable_args(tmp_path, body, **overrides): + """Args as they look when the YAML merge runs: already validated once.""" + base = _args("grpo", reward_key=None) + base.loss_type = "policy_loss" + base.custom_config_path = _write_yaml(tmp_path, body) + for key, value in overrides.items(): + setattr(base, key, value) + return base + + +def test_yaml_cannot_switch_to_an_estimator_that_needs_a_critic(arguments_module, tmp_path): + """Role composition and the offload flags were derived from the pre- + override `use_critic`, so accepting a YAML switch to PPO would leave the + run neither fully critic nor fully critic-free.""" + args = _overridable_args(tmp_path, "advantage_estimator: ppo\n") + + with pytest.raises(ValueError, match="critic setup"): + arguments_module.apply_custom_config_overrides(args) + + +# The three validators below were split out of `validate_algorithm_args` +# because the main path 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 move any value the missing three read. Each case here fails on the +# pre-fix code. + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + pytest.param("advantage_estimator: rloo\nkl_coef: 0.01\n", "nonzero --kl-coef", id="reward-side-kl"), + pytest.param( + "advantage_estimator: rloo\nnum_steps_per_rollout: 4\n", + "num-steps-per-rollout 1", + id="update-schedule", + ), + pytest.param( + "advantage_estimator: rloo\nglobal_batch_size: 999\n", + "one optimizer update per rollout", + id="batch-shape", + ), + ], +) +def test_yaml_cannot_bypass_the_late_running_validators(arguments_module, tmp_path, body, expected): + args = _overridable_args( + tmp_path, + body, + n_samples_per_prompt=8, + rollout_batch_size=16, + global_batch_size=128, + num_steps_per_rollout=None, + kl_coef=0.0, + max_staleness=0, + calculate_per_token_loss=True, + rewards_normalization=True, + normalize_advantages=False, + partial_rollout=False, + use_dynamic_global_batch_size=False, + hybrid=False, + fully_async=False, + ) + + with pytest.raises(ValueError, match=expected): + arguments_module.apply_custom_config_overrides(args) + + +def test_yaml_that_changes_the_update_schedule_gets_a_rederived_batch(arguments_module, tmp_path): + """Re-running the validator without its derivation rejected a legal config. + + `validate_batch_shape` reads `global_batch_size`, which the main path + derives from `num_steps_per_rollout` *before* the merge. A YAML switching + grpo@4-steps to rloo@1-step should get `rollout * n = 128`; the first + version of this fix compared against the stale 32 and refused it. + """ + args = _overridable_args( + tmp_path, + "advantage_estimator: rloo\nnum_steps_per_rollout: 1\n", + n_samples_per_prompt=8, + rollout_batch_size=16, + global_batch_size=32, + num_steps_per_rollout=4, + kl_coef=0.0, + max_staleness=0, + calculate_per_token_loss=True, + rewards_normalization=True, + normalize_advantages=False, + partial_rollout=False, + use_dynamic_global_batch_size=False, + hybrid=False, + fully_async=False, + ) + + arguments_module.apply_custom_config_overrides(args) + + assert args.global_batch_size == 128, "the merge should re-derive it, not keep the pre-merge value" + + +def test_yaml_without_algorithm_changes_is_accepted(arguments_module, tmp_path): + args = _overridable_args(tmp_path, "lr: 0.5\n") + arguments_module.apply_custom_config_overrides(args) + assert args.lr == 0.5 + assert args.advantage_estimator == "grpo" + + +def test_no_yaml_is_a_no_op(arguments_module): + args = _args("grpo", reward_key=None) + args.loss_type = "policy_loss" + args.custom_config_path = None + arguments_module.apply_custom_config_overrides(args) + + +def test_sft_runs_skip_the_algorithm_recheck(arguments_module, tmp_path): + """SFT never selects an estimator, so a stale one must not block it.""" + args = _overridable_args(tmp_path, "lr: 0.5\n", loss_type="sft", advantage_estimator="ppo") + arguments_module.apply_custom_config_overrides(args) + assert args.lr == 0.5 + + +def test_slime_validate_args_applies_overrides_through_the_helper(arguments_module): + """Guard the call site: the merge must go through the re-checking + helper.""" + import inspect + + src = inspect.getsource(arguments_module.slime_validate_args) + assert "apply_custom_config_overrides(args)" in src + assert "yaml.safe_load" not in src, "the YAML merge was inlined again, skipping the re-check" + + +def test_spec_with_an_unregistered_implementation_is_rejected_at_startup(arguments_module, monkeypatch): + """A registry typo must name itself, not KeyError inside a worker.""" + from dataclasses import replace + + from relax.algorithms.spec import ALGORITHM_SPECS + + broken = replace(ALGORITHM_SPECS["grpo"], advantage_fn="typo_does_not_exist") + monkeypatch.setitem(ALGORITHM_SPECS, "grpo", broken) + + with pytest.raises(ValueError, match="typo_does_not_exist"): + arguments_module.validate_algorithm_args(_args("grpo", reward_key=None)) + + +# ---------------- fully-async ---------------- + + +@pytest.mark.parametrize("estimator", ["grpo", "gspo", "sapo", "cispo"]) +def test_other_estimators_are_unaffected_by_fully_async(arguments_module, estimator): + args = _args(estimator, fully_async=True, reward_key=None) + arguments_module.validate_algorithm_args(args) diff --git a/tests/algorithms/test_dispatch_parity_vs_main.py b/tests/algorithms/test_dispatch_parity_vs_main.py new file mode 100644 index 000000000..302b7078e --- /dev/null +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -0,0 +1,524 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""The registry must route each algorithm exactly where main's if/elif did. + +Scope note, because it is easy to over-claim here. `relax/utils/training/ +ppo_utils.py` is byte-identical to main on this branch, so no estimator's or +policy loss's *maths* changed — every one of them still calls the same function +object it always did. Feeding both implementations the same tensors and +comparing outputs would therefore pass by construction and prove nothing. + +What the refactor did change is *routing*: which kernel each algorithm name +resolves to, and which capability flags gate the surrounding code. That is what +these tables pin down. They are transcribed from +main @ 4899b8f3a90489840a736897b4c341d87c6267cf, whose algorithm code last +changed in 98a72349c7d0368440eb6b0c6849e9d0f2ba8cef ("实现 RLOO advantage +estimator", #205) -- nothing between those two commits touched advantages.py, +loss.py, utils.py or ppo_utils.py: + + relax/components/advantages.py lines 176-218 (advantage if/elif) + relax/backends/megatron/loss.py lines 579-628 (the duplicate of it) + relax/backends/megatron/loss.py lines 691-694 (advantage normalisation) + relax/backends/megatron/loss.py line 894 (need_full_log_probs) + relax/backends/megatron/loss.py lines 851-854 (loss reducer) + relax/backends/megatron/loss.py line 953 (sequence-level KL) + relax/backends/megatron/loss.py lines 968-988 (policy loss if/elif) + relax/utils/utils.py lines 186,206 (reward normalisation) + relax/utils/arguments.py RLOO + REINFORCE++ startup constraints + +Do not regenerate these from the current implementation — that would turn the +comparison into the implementation checking itself. +""" + +import pytest + + +torch = pytest.importorskip("torch") + +from relax.algorithms import get_algorithm, list_algorithm_names # noqa: E402 +from relax.algorithms.advantages import ADVANTAGE_FNS # noqa: E402 +from relax.algorithms.policy import POLICY_LOSS_FNS # noqa: E402 +from relax.algorithms.rewards import REWARD_NORMALIZERS # noqa: E402 +from relax.utils.training import ppo_utils # noqa: E402 + + +MAIN_SHA = "4899b8f3a90489840a736897b4c341d87c6267cf" # the base this branch is rebased on + +# main advantages.py:176 — `if estimator in ["grpo", "gspo", "sapo", "cispo", "rloo"]` +# -> get_grpo_returns, etc. +MAIN_ADVANTAGE_KERNEL = { + "grpo": ppo_utils.get_grpo_returns, + "gspo": ppo_utils.get_grpo_returns, + "sapo": ppo_utils.get_grpo_returns, + "cispo": ppo_utils.get_grpo_returns, + "rloo": ppo_utils.get_grpo_returns, + "ppo": ppo_utils.get_advantages_and_returns_batch, + "reinforce_plus_plus": ppo_utils.get_reinforce_plus_plus_returns, + "reinforce_plus_plus_baseline": ppo_utils.get_reinforce_plus_plus_baseline_advantages, +} + +# main loss.py:968-988 +MAIN_POLICY_KERNEL = { + "grpo": ppo_utils.compute_policy_loss, + "gspo": ppo_utils.compute_policy_loss, + "sapo": ppo_utils.compute_sapo_loss, + "cispo": ppo_utils.compute_cispo_loss, + "rloo": ppo_utils.compute_rloo_loss, + "ppo": ppo_utils.compute_policy_loss, + "reinforce_plus_plus": ppo_utils.compute_policy_loss, + "reinforce_plus_plus_baseline": ppo_utils.compute_policy_loss, +} + +# main loss.py:953 — only gspo took the sequence-level KL branch. +MAIN_SEQUENCE_LEVEL_KL = {"gspo"} + +# main loss.py:894 — `args.use_opsm or estimator == "gspo"`. +MAIN_NEEDS_FULL_LOG_PROBS = {"gspo"} + +# main arguments.py:3319 — `use_critic = estimator == "ppo"`. +MAIN_NEEDS_CRITIC = {"ppo"} + +# main arguments.py:3164-3168 — reinforce_plus_plus{,_baseline} asserted normalize_advantages. +MAIN_REQUIRES_NORMALIZE_ADVANTAGES = {"reinforce_plus_plus", "reinforce_plus_plus_baseline"} + +# main utils.py:186 — group-mean whitelist; :206 — the subset that also divides by std. +# rloo is in the first (it is normalised per group) but not the second: its +# leave-one-out baseline deliberately keeps the reward scale. +MAIN_GROUP_NORMALIZED = {"grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline", "rloo"} +MAIN_GROUP_STD_NORMALIZED = {"grpo", "gspo", "sapo", "cispo"} + +# main loss.py:691-694 and 851-854 — the two duplicated REINFORCE++ name sets +# that drove `distributed_masked_normalize` and the mask-safe loss reducer. +MAIN_TOKEN_GLOBAL = {"reinforce_plus_plus", "reinforce_plus_plus_baseline"} + +# main arguments.py, the RLOO block and `_validate_reinforce_plus_plus_args`. +# Each of these was an `if args.advantage_estimator ...` on main and is a spec +# field here; the sets are what main actually enforced, not what the spec says. +MAIN_FORBIDS_NORMALIZE_ADVANTAGES = {"rloo"} +MAIN_REQUIRES_REWARDS_NORMALIZATION = {"rloo", "reinforce_plus_plus_baseline"} +MAIN_FORBIDS_REWARD_SIDE_KL = {"rloo", "reinforce_plus_plus_baseline"} +MAIN_REQUIRES_GLOBAL_TOKEN_LOSS = {"rloo"} +MAIN_REQUIRES_ON_POLICY_UPDATES = {"rloo"} +MAIN_MIN_GROUP_SIZE = {"rloo": 2, "reinforce_plus_plus_baseline": 2} + +MAIN_ALGORITHMS = sorted(MAIN_ADVANTAGE_KERNEL) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_advantage_routes_to_the_same_kernel_as_main(name): + """The estimator's maths is unchanged; only the lookup moved.""" + spec = get_algorithm(name) + fn = ADVANTAGE_FNS[spec.advantage_fn] + source = fn.__code__.co_names + expected = MAIN_ADVANTAGE_KERNEL[name].__name__ + assert expected in source, f"{name} no longer reaches {expected}; it calls {source}" + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_policy_loss_routes_to_the_same_kernel_as_main(name): + spec = get_algorithm(name) + fn = POLICY_LOSS_FNS[spec.policy_loss_fn] + expected = MAIN_POLICY_KERNEL[name].__name__ + assert expected in fn.__code__.co_names, f"{name} no longer reaches {expected}" + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_sequence_level_kl_matches_main(name): + assert (get_algorithm(name).kl_level == "sequence") is (name in MAIN_SEQUENCE_LEVEL_KL) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_needs_full_log_probs_matches_main(name): + assert get_algorithm(name).needs_full_log_probs is (name in MAIN_NEEDS_FULL_LOG_PROBS) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_needs_critic_matches_main(name): + assert get_algorithm(name).needs_critic is (name in MAIN_NEEDS_CRITIC) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_requires_normalize_advantages_matches_main(name): + spec = get_algorithm(name) + assert spec.requires_normalize_advantages is (name in MAIN_REQUIRES_NORMALIZE_ADVANTAGES) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_group_normalization_matches_main(name): + """utils.py had two overlapping whitelists; both are now spec fields.""" + normalizer = get_algorithm(name).reward_normalizer + assert (normalizer != "none") is (name in MAIN_GROUP_NORMALIZED) + assert (normalizer == "group_mean_std") is (name in MAIN_GROUP_STD_NORMALIZED) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_advantage_normalization_matches_main(name): + """main kept this as two identical name sets; drift between them was + silent.""" + assert (get_algorithm(name).advantage_normalization == "token_global") is (name in MAIN_TOKEN_GLOBAL) + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_startup_constraints_match_main(name): + spec = get_algorithm(name) + assert spec.forbids_normalize_advantages is (name in MAIN_FORBIDS_NORMALIZE_ADVANTAGES) + assert spec.requires_rewards_normalization is (name in MAIN_REQUIRES_REWARDS_NORMALIZATION) + assert spec.forbids_reward_side_kl is (name in MAIN_FORBIDS_REWARD_SIDE_KL) + assert spec.requires_global_token_loss is (name in MAIN_REQUIRES_GLOBAL_TOKEN_LOSS) + assert spec.requires_on_policy_updates is (name in MAIN_REQUIRES_ON_POLICY_UPDATES) + assert spec.min_group_size == MAIN_MIN_GROUP_SIZE.get(name, 1) + + +def test_no_algorithm_both_requires_and_forbids_advantage_normalization(): + """The two flags come from opposite sides of main's validation; a spec + setting both would make the algorithm unlaunchable in every + configuration.""" + for name in list_algorithm_names(): + spec = get_algorithm(name) + assert not (spec.requires_normalize_advantages and spec.forbids_normalize_advantages), name + + +def test_every_algorithm_main_supported_is_still_registered(): + """A migration that quietly dropped an algorithm would pass every other + test.""" + missing = set(MAIN_ALGORITHMS) - set(list_algorithm_names()) + assert not missing, f"{missing} were reachable on main {MAIN_SHA[:7]} and are gone now" + + +def test_reward_normalizer_identifiers_all_resolve(): + for name in list_algorithm_names(): + assert get_algorithm(name).reward_normalizer in REWARD_NORMALIZERS + + +# ---------------- the adapters must be identity wrappers ---------------- +# +# The co_names checks above only prove the kernel's name appears in the adapter's +# bytecode. They would still pass if the adapter scaled its input, dropped an +# argument, or threw the result away. These compare the adapter's output against +# calling the kernel directly with main's argument list, which is what actually +# pins "the wrapper adds nothing". + + +def _kl(lengths=(3, 2)): + return [torch.zeros(n, dtype=torch.float32) for n in lengths] + + +def _masks(lengths=(3, 2)): + return [torch.ones(n, dtype=torch.float32) for n in lengths] + + +def _args(estimator, **overrides): + from types import SimpleNamespace + + base = dict( + advantage_estimator=estimator, + kl_coef=0.0, + gamma=1.0, + lambd=1.0, + eps_clip=0.2, + eps_clip_high=0.3, + sapo_tau_pos=1.0, + sapo_tau_neg=1.05, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.mark.parametrize("name", ["grpo", "gspo", "sapo", "cispo"]) +def test_grpo_family_adapter_is_the_bare_kernel(name): + """main: torch.tensor(rewards, float32, device) then get_grpo_returns(...).""" + from relax.algorithms.advantages import compute_advantages_and_returns + + rewards, kl = [1.5, -2.0], _kl() + got, _ = compute_advantages_and_returns(_args(name), rewards=rewards, kl=kl, loss_masks=_masks()) + want = ppo_utils.get_grpo_returns(torch.tensor(rewards, dtype=torch.float32, device=kl[0].device), kl) + + assert len(got) == len(want) + for left, right in zip(got, want, strict=True): + assert torch.equal(left, right), name + + +def test_reinforce_plus_plus_baseline_adapter_is_the_bare_kernel(): + from relax.algorithms.advantages import compute_advantages_and_returns + + rewards, kl, masks = [3.0], [torch.tensor([2.0, 4.0])], [torch.ones(2)] + args = _args("reinforce_plus_plus_baseline", kl_coef=0.5) + + got, returns = compute_advantages_and_returns(args, rewards=rewards, kl=kl, loss_masks=masks) + want = ppo_utils.get_reinforce_plus_plus_baseline_advantages( + rewards=torch.tensor(rewards, dtype=torch.float32, device=kl[0].device), + kl=[torch.tensor([2.0, 4.0])], + loss_masks=masks, + ) + + for left, right in zip(got, want, strict=True): + assert torch.equal(left, right) + assert returns is got, "main aliased returns to advantages for this estimator" + + +def _loss_inputs(): + torch.manual_seed(20260726) + return torch.randn(8), torch.randn(8), torch.randn(8) + + +def test_ppo_clip_adapter_passes_mains_arguments(): + from relax.algorithms.policy import compute_policy_loss_for + + log_probs, ppo_kl, advantages = _loss_inputs() + args = _args("grpo") + got = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = ppo_utils.compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) + assert torch.equal(got[0], want[0]) and torch.equal(got[1], want[1]) + + +def test_sapo_adapter_passes_mains_arguments(): + from relax.algorithms.policy import compute_policy_loss_for + + log_probs, ppo_kl, advantages = _loss_inputs() + args = _args("sapo", sapo_tau_pos=1.3, sapo_tau_neg=1.7) + got = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = ppo_utils.compute_sapo_loss(ppo_kl=ppo_kl, advantages=advantages, tau_pos=1.3, tau_neg=1.7) + assert torch.equal(got[0], want[0]) and torch.equal(got[1], want[1]) + + +def test_sapo_adapter_uses_mains_defaults_when_args_omit_the_taus(): + from types import SimpleNamespace + + from relax.algorithms.policy import compute_policy_loss_for + + log_probs, ppo_kl, advantages = _loss_inputs() + bare = SimpleNamespace(advantage_estimator="sapo", eps_clip=0.2, eps_clip_high=0.3) + got = compute_policy_loss_for(bare, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = ppo_utils.compute_sapo_loss(ppo_kl=ppo_kl, advantages=advantages, tau_pos=1.0, tau_neg=1.05) + assert torch.equal(got[0], want[0]) + + +def test_cispo_adapter_passes_mains_arguments(): + """The one adapter taking four kernel arguments — most room to drop one.""" + from relax.algorithms.policy import compute_policy_loss_for + + log_probs, ppo_kl, advantages = _loss_inputs() + args = _args("cispo", eps_clip=0.15, eps_clip_high=9.0) + got = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = ppo_utils.compute_cispo_loss( + log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages, eps_clip=0.15, eps_clip_high=9.0 + ) + assert torch.equal(got[0], want[0]) and torch.equal(got[1], want[1]) + + +@pytest.fixture +def cp_disabled(monkeypatch): + """Minimal megatron.core.mpu so the reinforce++ kernel runs on CPU. + + get_reinforce_plus_plus_returns imports mpu inside the function and reads + only get_context_parallel_world_size(); at 1 it takes the non-gathering + branch, which is the configuration the rest of this file already assumes. + Stubbing it is what makes the adapter testable at all -- the alternative is + leaving the one selectable estimator with no numerical check, which is how + it got here. + """ + import sys + import types + + core = types.ModuleType("megatron.core") + core.mpu = types.SimpleNamespace( + get_context_parallel_world_size=lambda: 1, + get_context_parallel_rank=lambda: 0, + ) + megatron = types.ModuleType("megatron") + megatron.core = core + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + yield + + +def test_reinforce_plus_plus_adapter_is_the_bare_kernel(cp_disabled): + """The one live estimator whose adapter had only a co_names check. + + Coverage said it plainly: advantage_reinforce_plus_plus was never executed + by any test, so nothing would have caught the adapter dropping an argument + or reordering the keyword-only ones. + """ + from relax.algorithms.advantages import compute_advantages_and_returns + + rewards = [1.5, -2.0] + kl = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.4, 0.5])] + loss_masks = [torch.ones(3), torch.ones(2)] + response_lengths, total_lengths = [3, 2], [5, 4] + args = _args("reinforce_plus_plus", kl_coef=0.3, gamma=0.95) + + got, returns = compute_advantages_and_returns( + args, + rewards=rewards, + kl=kl, + loss_masks=loss_masks, + response_lengths=response_lengths, + total_lengths=total_lengths, + ) + want = ppo_utils.get_reinforce_plus_plus_returns( + rewards=torch.tensor(rewards, dtype=torch.float32, device=kl[0].device), + kl=kl, + loss_masks=loss_masks, + response_lengths=response_lengths, + total_lengths=total_lengths, + kl_coef=0.3, + gamma=0.95, + ) + + assert len(got) == len(want) + for left, right in zip(got, want, strict=True): + assert torch.equal(left, right) + assert returns is not got, "main copied the list before returning it" + + +def test_reinforce_plus_plus_adapter_forwards_gamma_and_kl_coef(cp_disabled): + """Both are read off args rather than passed through, so a swap or a + hardcoded default would survive the identity test above if it used the + kernel's defaults.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + inputs = dict( + rewards=[1.0, 1.0], + kl=[torch.tensor([0.5, 0.5]), torch.tensor([0.5])], + loss_masks=[torch.ones(2), torch.ones(1)], + response_lengths=[2, 1], + total_lengths=[3, 2], + ) + low, _ = compute_advantages_and_returns(_args("reinforce_plus_plus", kl_coef=0.0, gamma=1.0), **inputs) + high, _ = compute_advantages_and_returns(_args("reinforce_plus_plus", kl_coef=0.9, gamma=1.0), **inputs) + assert not torch.equal(low[0], high[0]), "kl_coef is not reaching the kernel" + + g_one, _ = compute_advantages_and_returns(_args("reinforce_plus_plus", kl_coef=0.0, gamma=1.0), **inputs) + g_half, _ = compute_advantages_and_returns(_args("reinforce_plus_plus", kl_coef=0.0, gamma=0.5), **inputs) + assert not torch.equal(g_one[0], g_half[0]), "gamma is not reaching the kernel" + + +# ---------------- RLOO: the new algorithm main added while this branch was open ---------------- +# +# RLOO arrived on main as inline code in three places. These compare the +# registry's versions against transcriptions of those places, so the migration +# is pinned numerically rather than only by which kernel name appears. + + +class _RlooSample: + """Only what `post_process_rewards` and `group_positions` read.""" + + def __init__(self, reward, group_index): + self.reward = reward + self.group_index = group_index + + def get_reward_value(self, args): + return self.reward + + +def _rloo_reward_args(n_samples_per_prompt): + from types import SimpleNamespace + + return SimpleNamespace( + advantage_estimator="rloo", + n_samples_per_prompt=n_samples_per_prompt, + rewards_normalization=True, + grpo_std_normalization=False, + custom_reward_post_process_path=None, + agentic_custom_advantage_path=None, + reward_key=None, + ) + + +def _main_rloo_normalized_rewards(args, samples, raw_rewards): + """Transcription of main utils.py:186-224, the `rloo` reward branch. + + Deliberately written out rather than imported: importing the production + helper would compare the implementation against itself. + """ + rewards = torch.tensor(raw_rewards, dtype=torch.float) + positions_by_group: dict[int, list[int]] = {} + for position, sample in enumerate(samples): + positions_by_group.setdefault(sample.group_index, []).append(position) + + normalized_rewards = torch.empty_like(rewards) + for positions in positions_by_group.values(): + group_rewards = rewards[positions] + group_size = group_rewards.shape[0] + mean_reward = group_rewards.mean() + scale = group_size / (group_size - 1) + normalized_rewards[positions] = scale * (group_rewards - mean_reward) + return normalized_rewards.tolist() + + +def test_rloo_reward_normalizer_matches_mains_inline_branch(): + from relax.utils.utils import post_process_rewards + + raw = [1.0, 0.0, 0.5, -2.0, 3.0, 3.0, 3.0, 0.25] + samples = [_RlooSample(value, group_index=index // 4) for index, value in enumerate(raw)] + args = _rloo_reward_args(4) + + got_raw, got_normalized = post_process_rewards(args, samples) + assert got_raw == raw, "main returned the raw rewards untouched alongside the normalised ones" + assert got_normalized == _main_rloo_normalized_rewards(args, samples, raw) + + +def test_rloo_normalizer_is_not_the_grpo_one(): + """Both are group-centred; only GRPO divides by the group std. + + Without this, routing `rloo` to `group_mean_std` by mistake would pass the + identity test above for any group whose std happens to be 1. + """ + from relax.algorithms.rewards import normalize_group_mean_std + from relax.utils.utils import post_process_rewards + + raw = [1.0, 0.0, 0.5, -2.0] + samples = [_RlooSample(value, group_index=0) for value in raw] + + grpo_args = _rloo_reward_args(4) + grpo_args.advantage_estimator = "grpo" + grpo_args.grpo_std_normalization = True + + _, rloo = post_process_rewards(_rloo_reward_args(4), samples) + assert rloo != normalize_group_mean_std(grpo_args, samples, raw) + + +def test_rloo_policy_loss_adapter_is_the_bare_kernel(): + """main loss.py:915-919 — `compute_rloo_loss(log_probs=..., + advantages=...)`.""" + from relax.algorithms.policy import compute_policy_loss_for + + torch.manual_seed(0) + log_probs, ppo_kl, advantages = torch.randn(8), torch.randn(8), torch.randn(8) + + got_loss, got_clipfrac = compute_policy_loss_for( + _args("rloo"), log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages + ) + want_loss, want_clipfrac = ppo_utils.compute_rloo_loss(log_probs=log_probs, advantages=advantages) + + assert torch.equal(got_loss, want_loss) + assert torch.equal(got_clipfrac, want_clipfrac) + + +def test_rloo_policy_loss_ignores_ppo_kl(): + """The unclipped objective has no ratio term; an adapter that quietly fed + `ppo_kl` to a clipped kernel would still match the kernel-name check.""" + from relax.algorithms.policy import compute_policy_loss_for + + torch.manual_seed(0) + log_probs, advantages = torch.randn(8), torch.randn(8) + first, _ = compute_policy_loss_for( + _args("rloo"), log_probs=log_probs, ppo_kl=torch.zeros(8), advantages=advantages + ) + second, _ = compute_policy_loss_for( + _args("rloo"), log_probs=log_probs, ppo_kl=torch.full((8,), 5.0), advantages=advantages + ) + assert torch.equal(first, second) + + +def test_rloo_advantage_adapter_is_the_grpo_broadcast(): + """main folded `rloo` into the `["grpo", "gspo", "sapo", "cispo"]` + branch.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + rewards, kl = [1.5, -2.0], _kl() + got, _ = compute_advantages_and_returns(_args("rloo"), rewards=rewards, kl=kl, loss_masks=_masks()) + want = ppo_utils.get_grpo_returns(torch.tensor(rewards, dtype=torch.float32, device=kl[0].device), kl) + + for left, right in zip(got, want, strict=True): + assert torch.equal(left, right) diff --git a/tests/algorithms/test_policy_loss_dispatch.py b/tests/algorithms/test_policy_loss_dispatch.py new file mode 100644 index 000000000..e0934178e --- /dev/null +++ b/tests/algorithms/test_policy_loss_dispatch.py @@ -0,0 +1,181 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Policy loss selection must come from the registry.""" + +import pathlib +import re +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + +from relax.algorithms.policy import POLICY_LOSS_FNS, compute_policy_loss_for # noqa: E402 +from relax.algorithms.spec import get_algorithm, list_algorithm_names # noqa: E402 +from relax.utils.training.ppo_utils import ( # noqa: E402 + compute_cispo_loss, + compute_policy_loss, + compute_sapo_loss, +) + + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +LOSS_PATH = REPO_ROOT / "relax" / "backends" / "megatron" / "loss.py" +SERVE_PATH = REPO_ROOT / "relax" / "components" / "advantages.py" + + +def _args(estimator, **overrides): + base = dict( + advantage_estimator=estimator, + eps_clip=0.2, + eps_clip_high=0.2, + sapo_tau_pos=1.0, + sapo_tau_neg=1.05, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def _tensors(): + torch.manual_seed(0) + return torch.randn(8), torch.randn(8), torch.randn(8) + + +# ---------------- registry ---------------- + + +def test_every_registered_loss_is_reachable_from_some_spec(): + """No dead entries: a loss no spec names can never be dispatched to. + + The inverse direction is `test_every_spec_policy_loss_id_is_registered`. + Together they pin the table to exactly what the registry uses, which is + what a hard-coded inventory of names was doing before -- except this + version fails for a reason instead of failing on every addition. + """ + referenced = {get_algorithm(name).policy_loss_fn for name in list_algorithm_names()} + assert set(POLICY_LOSS_FNS) == referenced + + +def test_every_spec_policy_loss_id_is_registered(): + for name in list_algorithm_names(): + assert get_algorithm(name).policy_loss_fn in POLICY_LOSS_FNS + + +# ---------------- adapters match their kernels ---------------- + + +def test_ppo_clip_matches_the_underlying_kernel(): + log_probs, ppo_kl, advantages = _tensors() + args = _args("grpo") + got = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) + assert torch.equal(got[0], want[0]) + assert torch.equal(got[1], want[1]) + + +def test_sapo_matches_the_underlying_kernel(): + log_probs, ppo_kl, advantages = _tensors() + got = compute_policy_loss_for(_args("sapo"), log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = compute_sapo_loss(ppo_kl=ppo_kl, advantages=advantages, tau_pos=1.0, tau_neg=1.05) + assert torch.equal(got[0], want[0]) + assert torch.equal(got[1], want[1]) + + +def test_cispo_matches_the_underlying_kernel(): + log_probs, ppo_kl, advantages = _tensors() + got = compute_policy_loss_for(_args("cispo"), log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = compute_cispo_loss( + log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages, eps_clip=0.2, eps_clip_high=0.2 + ) + assert torch.equal(got[0], want[0]) + assert torch.equal(got[1], want[1]) + + +def test_sapo_defaults_when_args_lack_tau_fields(): + log_probs, ppo_kl, advantages = _tensors() + args = SimpleNamespace(advantage_estimator="sapo", eps_clip=0.2, eps_clip_high=0.2) + got = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = compute_sapo_loss(ppo_kl=ppo_kl, advantages=advantages, tau_pos=1.0, tau_neg=1.05) + assert torch.equal(got[0], want[0]) + + +def test_sapo_taus_are_read_from_args(): + log_probs, ppo_kl, advantages = _tensors() + args = _args("sapo", sapo_tau_pos=2.0, sapo_tau_neg=3.0) + got = compute_policy_loss_for(args, log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + want = compute_sapo_loss(ppo_kl=ppo_kl, advantages=advantages, tau_pos=2.0, tau_neg=3.0) + assert torch.equal(got[0], want[0]) + + +@pytest.mark.parametrize("estimator", ["grpo", "gspo", "ppo", "reinforce_plus_plus"]) +def test_ppo_clip_family_share_one_loss(estimator): + log_probs, ppo_kl, advantages = _tensors() + reference = compute_policy_loss_for(_args("grpo"), log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + actual = compute_policy_loss_for(_args(estimator), log_probs=log_probs, ppo_kl=ppo_kl, advantages=advantages) + assert torch.equal(reference[0], actual[0]) + + +# ---------------- call sites no longer branch on names ---------------- + + +# Any comparison of the estimator against a literal name, in any of the shapes +# Python offers: `== "x"`, `!= "x"`, `in [...]`, `in {...}`, `in (...)`. +# +# This used to be a hand-written list of the exact strings the refactor had +# deleted. It missed a live one: the REINFORCE++ checks in loss.py were written +# `in {` and the list only banned `in [`, so two algorithm-name sets survived +# the refactor with a green test sitting on top of them. +ESTIMATOR_NAME_CHECK = re.compile(r"""advantage_estimator\s*(?:==|!=|\bin\b)\s*[\[{("']""") + + +def _name_checks(source: str) -> list[str]: + return [line.strip() for line in source.splitlines() if ESTIMATOR_NAME_CHECK.search(line)] + + +def test_loss_py_no_longer_branches_on_estimator_names(): + found = _name_checks(LOSS_PATH.read_text(encoding="utf-8")) + assert found == [], f"loss.py still compares the estimator to literal names: {found}" + + +def test_serve_path_no_longer_branches_on_estimator_names(): + found = _name_checks(SERVE_PATH.read_text(encoding="utf-8")) + assert found == [], f"components/advantages.py still compares the estimator to literal names: {found}" + + +def test_the_name_check_pattern_catches_every_spelling(): + """Guard the guard: the previous version of this test was blind to `in. + + {`. + """ + for spelling in ( + 'if args.advantage_estimator == "gspo":', + 'if args.advantage_estimator != "ppo":', + 'if args.advantage_estimator in ["grpo", "gspo"]:', + 'x = args.advantage_estimator in {"reinforce_plus_plus"}', + 'if self.config.advantage_estimator in ("ppo",):', + ): + assert _name_checks(spelling) == [spelling], spelling + for allowed in ( + "spec = get_algorithm(args.advantage_estimator)", + 'if get_algorithm(args.advantage_estimator).advantage_normalization == "token_global":', + ): + assert _name_checks(allowed) == [], allowed + + +def test_both_paths_delegate_to_the_shared_estimator(): + for path in (LOSS_PATH, SERVE_PATH): + src = path.read_text(encoding="utf-8") + assert "from relax.algorithms.advantages import" in src, f"{path.name} does not use the shared estimator" + + +def test_loss_py_reads_kl_level_and_full_log_probs_from_the_spec(): + src = LOSS_PATH.read_text(encoding="utf-8") + assert 'kl_level == "sequence"' in src + assert "needs_full_log_probs" in src + + +def test_neither_call_site_still_raises_not_implemented_for_estimators(): + for path in (LOSS_PATH, SERVE_PATH): + src = path.read_text(encoding="utf-8") + assert "advantage_estimator {" not in src, f"{path.name} still formats an estimator into NotImplementedError" diff --git a/tests/algorithms/test_post_process_rewards_dispatch.py b/tests/algorithms/test_post_process_rewards_dispatch.py new file mode 100644 index 000000000..64c5b6003 --- /dev/null +++ b/tests/algorithms/test_post_process_rewards_dispatch.py @@ -0,0 +1,169 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""post_process_rewards must dispatch through the registry, not an if/elif +chain.""" + +import inspect +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + +import relax.utils.utils as utils_mod # noqa: E402 + + +def _args(estimator="grpo", **overrides): + base = dict( + advantage_estimator=estimator, + n_samples_per_prompt=4, + rewards_normalization=True, + grpo_std_normalization=True, + custom_reward_post_process_path=None, + reward_key=None, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +class _Sample: + def __init__(self, group_index, reward): + self.group_index = group_index + self.reward = reward + + def get_reward_value(self, args): + return self.reward if not args.reward_key else self.reward[args.reward_key] + + +def test_source_has_no_algorithm_name_literals(): + """The estimator whitelists must be gone from post_process_rewards.""" + src = inspect.getsource(utils_mod.post_process_rewards) + for banned in ('"grpo"', '"gspo"', '"sapo"', '"cispo"', '"reinforce_plus_plus_baseline"'): + assert banned not in src, f"post_process_rewards still hardcodes {banned}" + + +def test_debug_subsample_source_has_no_algorithm_name_literals(): + src = inspect.getsource(utils_mod.get_debug_data) + for banned in ('"grpo"', '"gspo"', '"sapo"', '"cispo"', '"reinforce_plus_plus_baseline"'): + assert banned not in src, f"get_debug_data still hardcodes {banned}" + + +def test_returns_raw_and_normalized(): + args = _args("grpo") + samples = [_Sample(0, r) for r in (0.0, 1.0, 2.0, 3.0)] + raw, normalized = utils_mod.post_process_rewards(args, samples) + assert raw == [0.0, 1.0, 2.0, 3.0] + assert normalized != raw + assert abs(sum(normalized)) < 1e-5 + + +def test_identity_path_returns_raw_twice(): + args = _args("reinforce_plus_plus") + samples = [_Sample(0, r) for r in (0.0, 1.0, 2.0, 3.0)] + raw, normalized = utils_mod.post_process_rewards(args, samples) + assert normalized is raw + + +def test_rewards_normalization_off_returns_raw_twice(): + args = _args("grpo", rewards_normalization=False) + samples = [_Sample(0, r) for r in (0.0, 1.0, 2.0, 3.0)] + raw, normalized = utils_mod.post_process_rewards(args, samples) + assert normalized is raw + + +def test_baseline_estimator_centres_without_dividing_by_std(): + args = _args("reinforce_plus_plus_baseline") + samples = [_Sample(0, r) for r in (0.0, 1.0, 2.0, 3.0)] + _, normalized = utils_mod.post_process_rewards(args, samples) + assert normalized == pytest.approx([-1.5, -0.5, 0.5, 1.5]) + + +def test_custom_path_still_short_circuits(monkeypatch): + sentinel = (["raw"], ["norm"]) + monkeypatch.setattr(utils_mod, "load_function", lambda path: lambda a, s: sentinel) + args = _args("grpo", custom_reward_post_process_path="pkg.mod.fn") + assert utils_mod.post_process_rewards(args, []) is sentinel + + +def test_reward_key_selects_from_dict(): + args = _args("grpo", reward_key="score") + samples = [_Sample(0, {"score": r, "other": 99.0}) for r in (0.0, 1.0, 2.0, 3.0)] + raw, _ = utils_mod.post_process_rewards(args, samples) + assert raw == [0.0, 1.0, 2.0, 3.0] + + +def test_unknown_estimator_raises_from_the_registry(): + args = _args("not_an_algorithm") + samples = [_Sample(0, 1.0) for _ in range(4)] + with pytest.raises(KeyError, match="Unknown advantage estimator"): + utils_mod.post_process_rewards(args, samples) + + +# ---------------- raw_reward column stays scalar ---------------- + + +def _real_sample(group_index, reward, metadata=None): + """Use the production Sample so no field is accidentally missing.""" + from relax.utils.types import Sample + + return Sample( + group_index=group_index, + index=group_index, + tokens=[1, 2, 3], + response_length=2, + reward=reward, + metadata=metadata or {}, + ) + + +def _train_data_args(**overrides): + base = dict( + advantage_estimator="grpo", + n_samples_per_prompt=4, + rewards_normalization=True, + grpo_std_normalization=True, + custom_reward_post_process_path=None, + reward_key="score", + debug_train_only=True, # stop before dict_to_tensordict, keep plain lists + use_opd=False, + multimodal_keys=None, + use_rollout_routing_replay=False, + mask_offpolicy_in_partial_rollout=False, + ) + base.update(overrides) + return SimpleNamespace(**base) + + +def test_raw_reward_column_stays_scalar_when_metadata_overrides_some_samples(): + """Dict rewards plus a partial metadata override used to mix types. + + ``sample.reward`` is a dict for every run that uses --reward-key (and every + run using --reward-key), so falling back to it produced ``[0.8, {...}]``, + which blows up the TensorDict conversion. + """ + samples = [ + _real_sample(0, {"score": 1.0, "format": 0.5}, {"raw_reward": 0.8}), + _real_sample(0, {"score": 0.0, "format": 1.0}), + _real_sample(0, {"score": 1.0, "format": 0.0}), + _real_sample(0, {"score": 0.0, "format": 0.5}), + ] + + train_data = utils_mod.convert_samples_to_train_data(_train_data_args(), samples) + + assert train_data["raw_reward"] == [0.8, 0.0, 1.0, 0.0] + for value in train_data["raw_reward"]: + assert isinstance(value, float), f"{value!r} is not a scalar" + + +def test_raw_reward_column_untouched_without_metadata_overrides(): + samples = [_real_sample(0, {"score": float(i)}) for i in range(4)] + train_data = utils_mod.convert_samples_to_train_data(_train_data_args(), samples) + assert train_data["raw_reward"] == [0.0, 1.0, 2.0, 3.0] + + +def test_metadata_override_still_wins_for_scalar_rewards(): + """The original purpose of the override must keep working.""" + samples = [_real_sample(0, float(i), {"raw_reward": 9.0} if i == 0 else None) for i in range(4)] + train_data = utils_mod.convert_samples_to_train_data(_train_data_args(reward_key=None), samples) + assert train_data["raw_reward"] == [9.0, 1.0, 2.0, 3.0] diff --git a/tests/algorithms/test_reward_normalizers.py b/tests/algorithms/test_reward_normalizers.py new file mode 100644 index 000000000..df7b169af --- /dev/null +++ b/tests/algorithms/test_reward_normalizers.py @@ -0,0 +1,194 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Bit-exact characterization tests for the reward normalizers. + +``_legacy_post_process_rewards`` below is a frozen copy of the body of +``relax.utils.utils.post_process_rewards`` as of main@039ce87. It exists so +the refactor can be proven to change nothing: any difference in a single float +bit fails these tests. Do not "clean it up" — its value is that it is stale. +""" + +import random +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + +from relax.algorithms import get_algorithm # noqa: E402 +from relax.algorithms.rewards import REWARD_NORMALIZERS # noqa: E402 + + +_LEGACY_GROUP_NORM_ESTIMATORS = ["grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline"] +_LEGACY_STD_NORM_ESTIMATORS = ["grpo", "gspo", "sapo", "cispo"] + +ALL_ESTIMATORS = [ + "grpo", + "gspo", + "sapo", + "cispo", + "ppo", + "reinforce_plus_plus", + "reinforce_plus_plus_baseline", +] + +CONTIGUOUS_GROUPS = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2] +INTERLEAVED_GROUPS = [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2] + + +def _legacy_post_process_rewards(args, samples, raw_rewards): + """Frozen pre-refactor logic. + + Mirrors relax/utils/utils.py:181-207. + """ + if args.advantage_estimator in _LEGACY_GROUP_NORM_ESTIMATORS and args.rewards_normalization: + rewards = torch.tensor(raw_rewards, dtype=torch.float) + positions_by_group: dict[int, list[int]] = {} + for position, sample in enumerate(samples): + if sample.group_index is None: + raise ValueError("Sample.group_index is required for group reward normalization.") + if sample.group_index not in positions_by_group: + positions_by_group[sample.group_index] = [] + positions_by_group[sample.group_index].append(position) + + normalized_rewards = torch.empty_like(rewards) + for group_index, positions in positions_by_group.items(): + if len(positions) != args.n_samples_per_prompt: + raise ValueError( + f"Reward group {group_index} has {len(positions)} samples, expected {args.n_samples_per_prompt}." + ) + group_rewards = rewards[positions] + group_rewards = group_rewards - group_rewards.mean() + if args.advantage_estimator in _LEGACY_STD_NORM_ESTIMATORS and args.grpo_std_normalization: + group_rewards = group_rewards / (group_rewards.std() + 1e-6) + normalized_rewards[positions] = group_rewards + + return normalized_rewards.tolist() + + return raw_rewards + + +def _new_normalize(args, samples, raw_rewards): + """Mirrors the refactored dispatch in + relax.utils.utils.post_process_rewards.""" + if not args.rewards_normalization: + return raw_rewards + spec = get_algorithm(args.advantage_estimator) + return REWARD_NORMALIZERS[spec.reward_normalizer](args, samples, raw_rewards) + + +def _args(estimator, *, n=4, rewards_normalization=True, grpo_std_normalization=True): + return SimpleNamespace( + advantage_estimator=estimator, + n_samples_per_prompt=n, + rewards_normalization=rewards_normalization, + grpo_std_normalization=grpo_std_normalization, + ) + + +def _samples(group_indices): + return [SimpleNamespace(group_index=g) for g in group_indices] + + +def _assert_bitwise_equal(left, right): + left_t = torch.tensor(left, dtype=torch.float32) + right_t = torch.tensor(right, dtype=torch.float32) + assert left_t.shape == right_t.shape + assert torch.equal(left_t.view(torch.int32), right_t.view(torch.int32)), f"{left} != {right}" + + +def _reward_fixtures(): + rng = random.Random(20260725) + return { + "normal": [rng.uniform(-3, 3) for _ in range(12)], + "binary": [float(rng.randint(0, 1)) for _ in range(12)], + "all_equal": [0.7] * 12, + "one_group_collapsed": [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.5, 0.25, -1.0, 2.0, -0.5, 3.0], + "negatives": [-rng.uniform(0, 5) for _ in range(12)], + "duplicates": [1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 5.0, 5.0, 6.0, 6.0], + "large": [1e6, 1e6 + 1, 1e6 - 1, 1e6, 2e6, 2e6, 2e6, 2e6, 0.0, 1.0, 2.0, 3.0], + } + + +@pytest.mark.parametrize("estimator", ALL_ESTIMATORS) +@pytest.mark.parametrize("rewards_normalization", [True, False]) +@pytest.mark.parametrize("grpo_std_normalization", [True, False]) +@pytest.mark.parametrize("fixture_name", sorted(_reward_fixtures())) +@pytest.mark.parametrize("groups", [CONTIGUOUS_GROUPS, INTERLEAVED_GROUPS]) +def test_normalizer_is_bitwise_identical_to_legacy( + estimator, rewards_normalization, grpo_std_normalization, fixture_name, groups +): + raw = _reward_fixtures()[fixture_name] + args = _args( + estimator, + rewards_normalization=rewards_normalization, + grpo_std_normalization=grpo_std_normalization, + ) + samples = _samples(groups) + + expected = _legacy_post_process_rewards(args, samples, raw) + actual = _new_normalize(args, samples, raw) + + _assert_bitwise_equal(expected, actual) + + +def test_identity_normalizer_returns_the_same_list_object(): + """Non-normalising estimators must not copy — legacy returned raw_rewards + itself.""" + raw = [1.0, 2.0, 3.0, 4.0] + args = _args("reinforce_plus_plus") + samples = _samples([0, 0, 0, 0]) + assert _new_normalize(args, samples, raw) is raw + + +def test_missing_group_index_raises(): + args = _args("grpo") + samples = [SimpleNamespace(group_index=None) for _ in range(4)] + with pytest.raises(ValueError, match="group_index is required"): + _new_normalize(args, samples, [1.0, 2.0, 3.0, 4.0]) + + +def test_wrong_group_size_raises(): + args = _args("grpo", n=4) + samples = _samples([0, 0, 1, 1]) + with pytest.raises(ValueError, match="expected 4"): + _new_normalize(args, samples, [1.0, 2.0, 3.0, 4.0]) + + +def test_group_mean_normalizer_never_divides_by_std(): + """reinforce_plus_plus_baseline only centres, even with std normalisation + on.""" + args = _args("reinforce_plus_plus_baseline", n=4, grpo_std_normalization=True) + samples = _samples([0, 0, 0, 0]) + out = _new_normalize(args, samples, [0.0, 1.0, 2.0, 3.0]) + _assert_bitwise_equal(out, [-1.5, -0.5, 0.5, 1.5]) + + +def test_group_mean_std_respects_the_dr_grpo_switch(): + args_on = _args("grpo", n=4, grpo_std_normalization=True) + args_off = _args("grpo", n=4, grpo_std_normalization=False) + samples = _samples([0, 0, 0, 0]) + raw = [0.0, 1.0, 2.0, 3.0] + assert _new_normalize(args_on, samples, raw) != _new_normalize(args_off, samples, raw) + _assert_bitwise_equal(_new_normalize(args_off, samples, raw), [-1.5, -0.5, 0.5, 1.5]) + + +def test_grouping_is_driven_by_group_index_not_position(): + """Group membership follows ``group_index``, not batch position. + + Interleaving the same rewards into different groups therefore changes the + group statistics and so the normalized values -- which is what the + assertion below checks. (An earlier version of this docstring claimed the + opposite.) + """ + args = _args("grpo", n=4) + raw = [0.0, 1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 13.0] + + contiguous = _new_normalize(args, _samples([0, 0, 0, 0, 1, 1, 1, 1]), raw) + interleaved = _new_normalize(args, _samples([0, 1, 0, 1, 0, 1, 0, 1]), raw) + + assert sorted(round(v, 5) for v in contiguous) != sorted(round(v, 5) for v in interleaved) + # group 0 of the interleaved layout holds raw[0], raw[2], raw[4], raw[6] + group0 = [interleaved[i] for i in (0, 2, 4, 6)] + assert abs(sum(group0)) < 1e-5 From 434c22d4d1688b6db1f826874a7bd96c41936373 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Mon, 24 Aug 2026 20:43:12 +0800 Subject: [PATCH 2/6] refactor(algorithms): drive the critic capability from the registry, 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. --- relax/algorithms/__init__.py | 16 ++++- relax/algorithms/spec.py | 25 ++++++++ relax/backends/megatron/actor.py | 7 +- relax/components/critic.py | 10 +-- relax/core/controller.py | 3 +- relax/core/registry.py | 9 +-- relax/utils/training/data_fields.py | 10 +-- relax/utils/training/ppo_utils.py | 6 +- tests/algorithms/test_algos_roles.py | 95 ++++++++++++++++++++++++++++ 9 files changed, 161 insertions(+), 20 deletions(-) diff --git a/relax/algorithms/__init__.py b/relax/algorithms/__init__.py index 2d8c798bf..42a6e6431 100644 --- a/relax/algorithms/__init__.py +++ b/relax/algorithms/__init__.py @@ -9,7 +9,19 @@ the CPU-only test runner. """ -from relax.algorithms.spec import ALGORITHM_SPECS, AlgorithmSpec, get_algorithm, list_algorithm_names +from relax.algorithms.spec import ( + ALGORITHM_SPECS, + AlgorithmSpec, + algorithm_needs_critic, + get_algorithm, + list_algorithm_names, +) -__all__ = ["ALGORITHM_SPECS", "AlgorithmSpec", "get_algorithm", "list_algorithm_names"] +__all__ = [ + "ALGORITHM_SPECS", + "AlgorithmSpec", + "algorithm_needs_critic", + "get_algorithm", + "list_algorithm_names", +] diff --git a/relax/algorithms/spec.py b/relax/algorithms/spec.py index d124845db..17ccda9cb 100644 --- a/relax/algorithms/spec.py +++ b/relax/algorithms/spec.py @@ -221,6 +221,31 @@ def get_algorithm(name: str) -> AlgorithmSpec: raise KeyError(f"Unknown advantage estimator {name!r}. Available: {available}") from None +def algorithm_needs_critic(config) -> bool: + """Whether the configured algorithm runs a critic, read from the registry. + + The training pipeline asks this in five places -- role topology, the + critic's placement group, the device hand-off of the critic's ``values``, + the critic consumer's rollout fields, and the critic's own wait loop -- and + each of them used to compare ``advantage_estimator`` against ``"ppo"``. + That worked while PPO was the only value-based estimator, and silently + stopped working the moment the registry could accept a second one: + ``--advantage-estimator`` and ``ALGOS`` would take it, then none of the + value plumbing would switch on. + + ``args.use_critic`` carries the same answer, but only after + ``validate_algorithm_args`` has run; ``process_role`` and the controller's + placement logic read a config that may not have been through it yet. This + reads the spec directly so the answer does not depend on call order. + + Unknown or missing estimators answer False rather than raising: SFT and the + debug-only role paths reach these call sites with no estimator at all, and + an unknown name is rejected by argument parsing long before this matters. + """ + spec = ALGORITHM_SPECS.get(getattr(config, "advantage_estimator", None)) + return spec is not None and spec.needs_critic + + def list_algorithm_names() -> list[str]: """All registered algorithm names, in definition order.""" return list(ALGORITHM_SPECS) diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..a750e4507 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -26,6 +26,7 @@ from tensordict import TensorDict from transformers import AutoConfig, AutoTokenizer +from relax.algorithms import algorithm_needs_critic from relax.distributed.checkpoint_service.client.engine import create_client from relax.distributed.ray.train_actor import TrainRayActor from relax.engine.sft.eval.runner import run_sft_eval @@ -787,7 +788,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # and land on CPU (critic ``.cpu()`` s ``values`` before PUT). Inline # GAE + normalize_advantages need GPU tensors — dispatch here so the # rest of the pipeline can assume same-device inputs. - if self.args.advantage_estimator == "ppo": + if algorithm_needs_critic(self.args): cur_device = torch.cuda.current_device() for key in ("values", "loss_masks"): tensors = rollout_data.get(key) @@ -820,7 +821,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # advantages/returns via TransferQueue; every other path (including # PPO colocate) computes GAE inline from critic's ``values``. should_compute_gae_in_actor = self.args.compute_advantages_and_returns and not ( - self.args.advantage_estimator == "ppo" and self.args.fully_async and not self.args.hybrid + algorithm_needs_critic(self.args) and self.args.fully_async and not self.args.hybrid ) if should_compute_old_log_probs: @@ -2246,7 +2247,7 @@ def _put_data_to_transfer_queue(self, output_dict=None, batch_meta=None, rollout run(self.data_system_client.async_put(data=output_dict, metadata=batch_meta)) def _put_critic_values_to_transfer_queue(self, rollout_data: RolloutBatch) -> None: - if getattr(self.args, "advantage_estimator", None) != "ppo": + if not algorithm_needs_critic(self.args): return values = rollout_data.get("values") diff --git a/relax/components/critic.py b/relax/components/critic.py index ef01bbf5f..14b4883f7 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -120,23 +120,25 @@ def _wait_for_rollout_data(self) -> None: self._rollout_barrier.wait_offloaded_sync() def train(self) -> None: - is_ppo = getattr(self.config, "advantage_estimator", None) == "ppo" + from relax.algorithms import algorithm_needs_critic + + has_critic = algorithm_needs_critic(self.config) while self.step < self.config.num_rollout: - if is_ppo: + if has_critic: self._wait_for_rollout_data() # In PPO colocate the actor waits for ``self.step`` to advance # past the current round before waking up, so block on training # completion here. Non-PPO critic is not on any live service graph # and keeps the historical fire-and-forget. train_ref = self.critic_model.async_train(self.step) - if is_ppo: + if has_critic: ray.get(train_ref) # Note: save_model runs inside ``train_critic`` (backend) while the # model is still awake, so no explicit save call here. # In critic-only warmup, actor+advantages never consume the partition, # so critic must clear it itself; steady-state clearing stays with actor. - if is_ppo and self.step < getattr(self.config, "num_critic_only_steps", 0): + if has_critic and self.step < getattr(self.config, "num_critic_only_steps", 0): run(self.data_system_client.async_clear_partition(partition_id=f"train_{self.step}")) try: diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..bf78dda14 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -30,6 +30,7 @@ deploy_agentic_chat_api_services, shutdown_agentic_chat_api_services, ) +from relax.algorithms import algorithm_needs_critic from relax.core.optional_roles import register_extra_roles from relax.core.registry import ALGOS, ROLES, process_role from relax.core.service import Service, create_placement_group @@ -104,7 +105,7 @@ def _actor_rollout_pg_roles(config: Namespace) -> list[str]: matches the actor's; otherwise critic runs on its own placement group. """ roles = list(ACTOR_ROLLOUT_PG_ROLES) - if getattr(config, "advantage_estimator", None) != "ppo": + if not algorithm_needs_critic(config): return roles resource = getattr(config, "resource", None) or {} if resource.get("critic") == resource.get("actor"): diff --git a/relax/core/registry.py b/relax/core/registry.py index ce59c24d8..e70c0316e 100644 --- a/relax/core/registry.py +++ b/relax/core/registry.py @@ -11,7 +11,7 @@ def __str__(self) -> str: return self.value -from relax.algorithms import ALGORITHM_SPECS +from relax.algorithms import ALGORITHM_SPECS, algorithm_needs_critic from relax.components.actor import Actor from relax.components.actor_fwd import ActorFwd from relax.components.advantages import Advantages @@ -91,8 +91,9 @@ def _rl_roles(*, needs_critic: bool) -> dict: service-registration time, because both sides read the same dict. This decides which class a role maps to, *not* which roles the controller - walks -- that is ``process_role``'s job and it is deliberately left - untouched. ``controller.py`` iterates ``list(process_role(config))`` and + walks -- that is ``process_role``'s job, which reads the same + ``needs_critic`` through ``algorithm_needs_critic``. ``controller.py`` + iterates ``list(process_role(config))`` and skips any role missing from this dict, so an algorithm without a critic simply never matches the ``critic`` member the role sets already carry. """ @@ -125,7 +126,7 @@ def process_role(config): return ROLES_TRAIN_ONLY if getattr(config, "loss_type", None) == "sft": return ROLES_SFT_ONLY - if getattr(config, "advantage_estimator", None) == "ppo": + if algorithm_needs_critic(config): if config.fully_async: if getattr(config, "true_on_policy_mode", False): return ROLES_PPO_FULLY_ASYNC_ON_POLICY diff --git a/relax/utils/training/data_fields.py b/relax/utils/training/data_fields.py index ebd6479af..b7aa888d3 100644 --- a/relax/utils/training/data_fields.py +++ b/relax/utils/training/data_fields.py @@ -32,12 +32,14 @@ def build_data_fields(args: Namespace, *, consumer: str = "actor") -> list[str]: fields.append("multimodal_train_inputs") return fields - is_ppo = getattr(args, "advantage_estimator", None) == "ppo" + from relax.algorithms import algorithm_needs_critic - if is_ppo and consumer == "critic": + has_critic = algorithm_needs_critic(args) + + if has_critic and consumer == "critic": return _base_rollout_fields(args) - if is_ppo and consumer == "advantages": + if has_critic and consumer == "advantages": # PPO colocate never runs actor_fwd, so ref/log_probs are only # requested when kl_coef != 0 (i.e. an actor_fwd role is present). fields = _base_rollout_fields(args) @@ -53,7 +55,7 @@ def build_data_fields(args: Namespace, *, consumer: str = "actor") -> list[str]: return fields fields = _base_rollout_fields(args) - if is_ppo: + if has_critic: # PPO colocate: actor consumes critic's ``values`` and computes GAE # inline. Fully_async: standalone Advantages service produces # ``advantages``/``returns``, actor just pulls the finished tensors. diff --git a/relax/utils/training/ppo_utils.py b/relax/utils/training/ppo_utils.py index e4eab1310..b5bd4790f 100644 --- a/relax/utils/training/ppo_utils.py +++ b/relax/utils/training/ppo_utils.py @@ -11,6 +11,7 @@ import torch.distributed as dist import torch.nn.functional as F +from relax.algorithms import algorithm_needs_critic from relax.utils.logging_utils import get_logger @@ -18,12 +19,13 @@ def validate_ppo_config(config: Namespace) -> None: - if getattr(config, "advantage_estimator", None) != "ppo": + if not algorithm_needs_critic(config): return resource = getattr(config, "resource", None) or {} if "critic" not in resource: - raise ValueError("--advantage-estimator ppo requires a 'critic' entry in --resource.") + estimator = getattr(config, "advantage_estimator", None) + raise ValueError(f"--advantage-estimator {estimator} requires a 'critic' entry in --resource.") if getattr(config, "fully_async", False) or getattr(config, "hybrid", False): raise ValueError("PPO does not currently support --fully-async or --hybrid.") diff --git a/tests/algorithms/test_algos_roles.py b/tests/algorithms/test_algos_roles.py index d5320a24b..eeff2a0fe 100644 --- a/tests/algorithms/test_algos_roles.py +++ b/tests/algorithms/test_algos_roles.py @@ -133,3 +133,98 @@ def test_each_algorithm_gets_an_independent_role_dict(): names = list_algorithm_names() for left, right in zip(names, names[1:], strict=False): assert ALGOS[left] is not ALGOS[right], f"{left} and {right} share one dict object" + + +# ---------------- a second critic algorithm must need no new name checks ---------------- + + +def _register_second_critic_algorithm(monkeypatch): + """Add a value-based estimator to the registry for the duration of a test. + + Returns its name. Uses PPO's own spec so the only thing under test is + whether the pipeline keys off `needs_critic` or off the literal `"ppo"`. + """ + import dataclasses + + from relax.algorithms import ALGORITHM_SPECS + from relax.algorithms import spec as spec_module + + name = "ppo_second" + clone = dataclasses.replace(ALGORITHM_SPECS["ppo"], name=name) + monkeypatch.setitem(spec_module.ALGORITHM_SPECS, name, clone) + assert ALGORITHM_SPECS[name].needs_critic is True + return name + + +def _args_for(name, **overrides): + from argparse import Namespace + + base = dict( + advantage_estimator=name, + multimodal_keys=None, + kl_coef=0.0, + fully_async=False, + hybrid=False, + use_opd=False, + use_rollout_logprobs=False, + true_on_policy_mode=False, + debug_rollout_only=False, + debug_train_only=False, + loss_type=None, + ) + base.update(overrides) + return Namespace(**base) + + +def test_a_second_critic_algorithm_gets_the_critic_rollout_fields(monkeypatch): + """`values` must reach the advantages consumer without naming the algorithm. + + This is the failure the registry was supposed to make impossible: argparse + and `ALGOS` accept a second value-based estimator, and then the value + plumbing silently does not switch on because it compares against `"ppo"`. + """ + from relax.utils.training.data_fields import build_data_fields + + name = _register_second_critic_algorithm(monkeypatch) + + ppo_fields = build_data_fields(_args_for("ppo"), consumer="advantages") + new_fields = build_data_fields(_args_for(name), consumer="advantages") + assert new_fields == ppo_fields, "a second critic algorithm sees different fields than PPO" + assert "values" in new_fields + + # and the critic consumer's own set, which is the base set rather than the actor's + assert build_data_fields(_args_for(name), consumer="critic") == build_data_fields( + _args_for("ppo"), consumer="critic" + ) + + +def test_a_second_critic_algorithm_is_told_it_needs_a_critic_resource(monkeypatch): + """Startup validation is keyed on the capability, not on the name.""" + import pytest as _pytest + + from relax.utils.training.ppo_utils import validate_ppo_config + + name = _register_second_critic_algorithm(monkeypatch) + + with _pytest.raises(ValueError, match="requires a 'critic' entry"): + validate_ppo_config(_args_for(name, resource={"actor": "a"})) + + # a non-critic estimator is still waved through + validate_ppo_config(_args_for("grpo", resource={"actor": "a"})) + + +@requires_megatron +def test_a_second_critic_algorithm_walks_the_critic_role_topology(monkeypatch): + """`process_role` decides which roles the controller walks at all.""" + from relax.core.registry import process_role + + name = _register_second_critic_algorithm(monkeypatch) + + # identity, not membership: every role set carries a `critic` member and + # `ALGOS` is what filters it out, so comparing member names would pass even + # if the second algorithm fell through to the non-critic topology. + assert process_role(_args_for(name)) is process_role(_args_for("ppo")) + assert process_role(_args_for(name)) is not process_role(_args_for("grpo")) + + # and the fully-async split follows too, rather than only the colocate one + assert process_role(_args_for(name, fully_async=True)) is process_role(_args_for("ppo", fully_async=True)) From 8c56fbe95faeb5f2e72044eaf4851c889a1ae429 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Mon, 24 Aug 2026 20:43:27 +0800 Subject: [PATCH 3/6] fix(arguments): keep a YAML global_batch_size instead of deriving over 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. --- relax/utils/arguments.py | 17 ++++++ .../algorithms/test_arguments_spec_driven.py | 56 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 64f9dc33d..c7313e54b 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -3200,7 +3200,24 @@ def apply_custom_config_overrides(args) -> None: # `rollout * n`, and the validator instead saw the value derived from 4. # `enforce_consistency=False` because the stale value is, by construction, # the one derived before the merge -- comparing against it is the bug. + # + # A YAML file that *names* `global_batch_size` is a different case: that is + # not a stale value left over from an earlier derivation, it is the override + # this function exists to apply. Re-deriving over it wrote the YAML's value + # and then replaced it in the next statement, so the run used neither the + # configured number nor an error -- the one outcome the override contract + # rules out. Derive first so the comparison has something to name, then + # refuse the conflict rather than picking a winner. + yaml_global_batch_size = data.get("global_batch_size") derive_global_batch_size(args, enforce_consistency=False) + if yaml_global_batch_size is not None and args.global_batch_size != yaml_global_batch_size: + raise ValueError( + f"--custom-config-path sets global_batch_size to {yaml_global_batch_size}, but " + f"num_steps_per_rollout {args.num_steps_per_rollout} over rollout_batch_size " + f"{args.rollout_batch_size} * n_samples_per_prompt {args.n_samples_per_prompt} derives " + f"{args.global_batch_size}. Remove one of the two from the YAML -- whichever you drop, " + f"the other is what the run would otherwise have used without saying so." + ) validate_batch_shape(args) if args.use_critic != use_critic_before_override: # Role composition and the offload flags were derived from the pre-override diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py index 7c5eec96d..ab6798eb4 100644 --- a/tests/algorithms/test_arguments_spec_driven.py +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -314,3 +314,59 @@ def test_spec_with_an_unregistered_implementation_is_rejected_at_startup(argumen def test_other_estimators_are_unaffected_by_fully_async(arguments_module, estimator): args = _args(estimator, fully_async=True, reward_key=None) arguments_module.validate_algorithm_args(args) + + +# ---------------- a YAML global_batch_size must not be derived over ---------------- + + +def _batch_args(tmp_path, body, **overrides): + """Args shaped for the batch-size derivation, already validated once.""" + base = _overridable_args(tmp_path, body, **overrides) + base.rollout_batch_size = 32 + base.n_samples_per_prompt = 4 + base.num_steps_per_rollout = 4 + base.global_batch_size = 32 # 32 * 4 // 4, i.e. what the pre-merge derivation wrote + base.micro_batch_size = 1 + base.use_dynamic_batch_size = False + for key, value in overrides.items(): + setattr(base, key, value) + return base + + +def test_yaml_global_batch_size_conflicting_with_the_derivation_is_refused(arguments_module, tmp_path): + """The override used to be written and then silently replaced. + + `apply_custom_config_overrides` merges the YAML, then re-derives + `global_batch_size` from `num_steps_per_rollout`. With the derivation + unconditional, a YAML that names `global_batch_size` had its value assigned + by the merge loop and overwritten one statement later -- so the run used + neither the configured number nor an error, which is the single outcome the + "YAML key overrides the argument" contract does not allow. + """ + args = _batch_args(tmp_path, "global_batch_size: 999\n") + + with pytest.raises(ValueError, match="sets global_batch_size to 999"): + arguments_module.apply_custom_config_overrides(args) + + +def test_yaml_global_batch_size_agreeing_with_the_derivation_survives(arguments_module, tmp_path): + """Naming the value the derivation would reach anyway is not a conflict.""" + args = _batch_args(tmp_path, "global_batch_size: 32\n") + + arguments_module.apply_custom_config_overrides(args) + + assert args.global_batch_size == 32 + + +def test_yaml_that_only_moves_a_derivation_input_still_re_derives(arguments_module, tmp_path): + """The behaviour the `enforce_consistency=False` call was added for. + + Switching `num_steps_per_rollout` from 4 to 1 has to produce + `rollout * n`; the pre-merge value of 32 is stale by construction and must + not be compared against. + """ + args = _batch_args(tmp_path, "num_steps_per_rollout: 1\n") + + arguments_module.apply_custom_config_overrides(args) + + assert args.global_batch_size == 128 # 32 * 4 // 1 From 71e976528fe524046894f828ab7049654f02230a Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Mon, 24 Aug 2026 20:43:47 +0800 Subject: [PATCH 4/6] test(algorithms): compare the GAE adapter against main's numbers, not 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. --- .../test_dispatch_parity_vs_main.py | 119 ++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/tests/algorithms/test_dispatch_parity_vs_main.py b/tests/algorithms/test_dispatch_parity_vs_main.py index 302b7078e..0faf3de21 100644 --- a/tests/algorithms/test_dispatch_parity_vs_main.py +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -522,3 +522,122 @@ def test_rloo_advantage_adapter_is_the_grpo_broadcast(): for left, right in zip(got, want, strict=True): assert torch.equal(left, right) + + +# ---------------- GAE: the adapter main's two call sites disagreed on ---------------- +# +# `advantage_gae` had only a co_names check, 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`, which one of main's two call sites passed and the +# other did not. A dropped or reordered argument here does not raise -- it reads +# the wrong token positions and trains on them. + + +def _gae_inputs(): + """Fresh tensors per call: `advantage_gae` mutates `kl` in place (`k *= ...`).""" + kl = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.4, 0.5])] + values = [torch.tensor([0.5, 0.25, 0.125]), torch.tensor([1.0, 2.0])] + return dict( + rewards=[1.5, -2.0], + kl=kl, + values=values, + response_lengths=[3, 2], + total_lengths=[3, 2], + ) + + +def _main_gae(kl_coef, gamma, lambd, *, padded_total_lengths=None): + """main's PPO branch, transcribed, not regenerated. + + components/advantages.py:181-193 and the megatron duplicate at + loss.py:585-602. The only difference between the two is that the megatron + one forwards `padded_total_lengths`; both shape the reward identically. + """ + inputs = _gae_inputs() + old_rewards, kl = inputs["rewards"], inputs["kl"] + rewards = [] + for reward, k in zip(old_rewards, kl, strict=False): + k *= -kl_coef + cp_rank = 0 # the fixture pins mpu.get_context_parallel_rank() to 0 + if cp_rank == 0: + k[-1] += reward + rewards.append(k) + return ppo_utils.get_advantages_and_returns_batch( + inputs["total_lengths"], + inputs["response_lengths"], + inputs["values"], + rewards, + gamma, + lambd, + padded_total_lengths=padded_total_lengths, + ) + + +@pytest.mark.parametrize( + "kl_coef,gamma,lambd", + [ + (0.0, 1.0, 1.0), # the degenerate case the co_names check implied was enough + (0.05, 0.99, 0.95), # non-zero kl_coef and real discounting + ], +) +def test_gae_adapter_matches_mains_numbers(cp_disabled, kl_coef, gamma, lambd): + """Every element, not just the kernel's name in the bytecode.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + args = _args("ppo", kl_coef=kl_coef, gamma=gamma, lambd=lambd) + got_adv, got_ret = compute_advantages_and_returns(args, **_gae_inputs()) + want_adv, want_ret = _main_gae(kl_coef, gamma, lambd) + + for got, want, label in ((got_adv, want_adv, "advantages"), (got_ret, want_ret, "returns")): + assert len(got) == len(want), label + for left, right in zip(got, want, strict=True): + torch.testing.assert_close(left, right, rtol=0, atol=0, msg=label) + + +def test_gae_terminal_reward_lands_on_the_last_token(cp_disabled): + """`k[-1] += reward` is the whole reward signal; dropping it trains on KL alone.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + args = _args("ppo", kl_coef=0.0, gamma=1.0, lambd=1.0) + with_reward, _ = compute_advantages_and_returns(args, **_gae_inputs()) + + zeroed = _gae_inputs() + zeroed["rewards"] = [0.0, 0.0] + without_reward, _ = compute_advantages_and_returns(args, **zeroed) + + assert not torch.equal(with_reward[0], without_reward[0]), "the terminal reward changed nothing" + + +def test_gae_adapter_forwards_padded_total_lengths_to_the_kernel(cp_disabled, monkeypatch): + """The argument the two call sites disagreed on must survive the adapter. + + Scope, stated because it is easy to over-read: `padded_total_lengths` is + only *consumed* when `cp_size > 1` (ppo_utils.py, the `all_gather_with_cp` + branch), and this file runs at cp_size 1. So this pins that the adapter + hands the value through unchanged and in the right keyword -- the "dropped + or reordered argument" failure -- not that the padded slicing itself is + correct. That needs a real context-parallel group and is not covered here. + """ + from relax.algorithms import advantages as advantages_module + + seen = {} + + def spy(*args, **kwargs): + seen["args"] = args + seen["kwargs"] = kwargs + return ([torch.zeros(3)], [torch.zeros(3)]) + + monkeypatch.setattr(advantages_module, "get_advantages_and_returns_batch", spy) + + padded = [8, 8] + advantages_module.compute_advantages_and_returns( + _args("ppo", kl_coef=0.0, gamma=1.0, lambd=1.0), + **_gae_inputs(), + padded_total_lengths=padded, + ) + + assert seen["kwargs"].get("padded_total_lengths") == padded + # main's positional order: total_lengths, response_lengths, values, rewards, gamma, lambd + assert seen["args"][0] == [3, 2] + assert seen["args"][1] == [3, 2] + assert seen["args"][4] == 1.0 and seen["args"][5] == 1.0 From 45d7d467772dcfc832d7baa72e6f3c0bb54649d6 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Mon, 24 Aug 2026 20:44:02 +0800 Subject: [PATCH 5/6] fix(algorithms): reject an unsupported kl_level or advantage_normalization `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. --- relax/algorithms/spec.py | 33 +++++++++++++++ tests/algorithms/test_algorithm_registry.py | 45 +++++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/relax/algorithms/spec.py b/relax/algorithms/spec.py index 17ccda9cb..6b6a931f8 100644 --- a/relax/algorithms/spec.py +++ b/relax/algorithms/spec.py @@ -19,6 +19,16 @@ from dataclasses import dataclass +# The two enum-like fields below are consumed by equality checks in +# `relax/backends/megatron/loss.py` -- `advantage_normalization == "token_global"` +# at 659 and 819, `kl_level == "sequence"` at 919. Anything that is not the +# awaited string takes the *other* branch, so a typo in a registry entry does +# not fail, it silently selects a different formula. These sets are what +# `__post_init__` checks against. +KL_LEVELS = frozenset({"token", "sequence"}) +ADVANTAGE_NORMALIZATIONS = frozenset({"whiten", "token_global"}) + + @dataclass(frozen=True) class AlgorithmSpec: """Everything the training pipeline needs to know about one algorithm.""" @@ -121,6 +131,29 @@ class AlgorithmSpec: only sees a slice of). Those get their own field; do not fold them here. """ + def __post_init__(self) -> None: + """Reject an unsupported enum value while the registry is being built. + + ``ALGORITHM_SPECS`` is a module-level literal, so this runs at import: + a typo cannot reach a worker, let alone a training step. The + implementation identifiers are already resolved eagerly for the same + reason (``_assert_spec_implementations_resolve`` in ``arguments.py``); + these two fields were the ones left unchecked, and they are the ones + whose failure is silent rather than loud -- a bad ``advantage_fn`` + raises a KeyError, a bad ``kl_level`` just trains with token-level KL + and never says so. + """ + for field, value, allowed in ( + ("kl_level", self.kl_level, KL_LEVELS), + ("advantage_normalization", self.advantage_normalization, ADVANTAGE_NORMALIZATIONS), + ): + if value not in allowed: + raise ValueError( + f"AlgorithmSpec({self.name!r}) has {field}={value!r}, which no call site matches; " + f"the run would silently take the default branch instead. " + f"Expected one of {sorted(allowed)}." + ) + @property def is_group_normalized(self) -> bool: """Whether rewards get normalised per prompt group on the rollout diff --git a/tests/algorithms/test_algorithm_registry.py b/tests/algorithms/test_algorithm_registry.py index 5e8706385..9ff313675 100644 --- a/tests/algorithms/test_algorithm_registry.py +++ b/tests/algorithms/test_algorithm_registry.py @@ -153,3 +153,48 @@ def test_every_spec_declares_a_kl_level_the_loss_knows_how_to_read(): """ for name in list_algorithm_names(): assert get_algorithm(name).kl_level in ("token", "sequence"), name + + +# ---------------- enum-like fields must fail loudly, not silently ---------------- + + +@pytest.mark.parametrize( + "field,bad", + [ + ("kl_level", "Sequence"), # right word, wrong case + ("kl_level", "seq"), + ("advantage_normalization", "token-global"), # hyphen instead of underscore + ("advantage_normalization", "none"), + ], +) +def test_an_unsupported_enum_value_is_refused_when_the_spec_is_built(field, bad): + """These two fields are compared for equality, so a typo picks a formula. + + `loss.py` asks `advantage_normalization == "token_global"` and + `kl_level == "sequence"`; every other string takes the else branch. Unlike + `advantage_fn`, which blows up with a KeyError the first time it is looked + up, a misspelled value here starts training successfully and quietly uses + the wrong KL level or the wrong advantage normalisation. The registry is a + module-level literal, so validating in `__post_init__` moves that from a + silent wrong-maths run to an import-time error. + """ + from relax.algorithms.spec import AlgorithmSpec + + kwargs = dict( + name="probe", + reward_normalizer="none", + advantage_fn="grpo_broadcast", + policy_loss_fn="ppo_clip", + ) + kwargs[field] = bad + + with pytest.raises(ValueError, match=field): + AlgorithmSpec(**kwargs) + + +def test_the_supported_values_are_the_ones_the_shipped_specs_use(): + """Guards against the allow-list drifting away from the registry it guards.""" + from relax.algorithms.spec import ADVANTAGE_NORMALIZATIONS, ALGORITHM_SPECS, KL_LEVELS + + assert {s.kl_level for s in ALGORITHM_SPECS.values()} <= KL_LEVELS + assert {s.advantage_normalization for s in ALGORITHM_SPECS.values()} <= ADVANTAGE_NORMALIZATIONS From bae4ced39d82d8bfe74fb9936b64eb0eb061e2ef Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Mon, 24 Aug 2026 20:51:37 +0800 Subject: [PATCH 6/6] style(algorithms): keep the new docstring summaries on one line `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. --- tests/algorithms/test_algorithm_registry.py | 4 ++-- tests/algorithms/test_algos_roles.py | 2 +- tests/algorithms/test_dispatch_parity_vs_main.py | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/algorithms/test_algorithm_registry.py b/tests/algorithms/test_algorithm_registry.py index 9ff313675..116c0d19a 100644 --- a/tests/algorithms/test_algorithm_registry.py +++ b/tests/algorithms/test_algorithm_registry.py @@ -170,7 +170,7 @@ def test_every_spec_declares_a_kl_level_the_loss_knows_how_to_read(): def test_an_unsupported_enum_value_is_refused_when_the_spec_is_built(field, bad): """These two fields are compared for equality, so a typo picks a formula. - `loss.py` asks `advantage_normalization == "token_global"` and + `loss.py` asks for `advantage_normalization == "token_global"`, and for `kl_level == "sequence"`; every other string takes the else branch. Unlike `advantage_fn`, which blows up with a KeyError the first time it is looked up, a misspelled value here starts training successfully and quietly uses @@ -193,7 +193,7 @@ def test_an_unsupported_enum_value_is_refused_when_the_spec_is_built(field, bad) def test_the_supported_values_are_the_ones_the_shipped_specs_use(): - """Guards against the allow-list drifting away from the registry it guards.""" + """The allow-list must not drift from the registry it guards.""" from relax.algorithms.spec import ADVANTAGE_NORMALIZATIONS, ALGORITHM_SPECS, KL_LEVELS assert {s.kl_level for s in ALGORITHM_SPECS.values()} <= KL_LEVELS diff --git a/tests/algorithms/test_algos_roles.py b/tests/algorithms/test_algos_roles.py index eeff2a0fe..47e2e5b7c 100644 --- a/tests/algorithms/test_algos_roles.py +++ b/tests/algorithms/test_algos_roles.py @@ -177,7 +177,7 @@ def _args_for(name, **overrides): def test_a_second_critic_algorithm_gets_the_critic_rollout_fields(monkeypatch): - """`values` must reach the advantages consumer without naming the algorithm. + """`values` must reach the advantages consumer, unprompted by any name. This is the failure the registry was supposed to make impossible: argparse and `ALGOS` accept a second value-based estimator, and then the value diff --git a/tests/algorithms/test_dispatch_parity_vs_main.py b/tests/algorithms/test_dispatch_parity_vs_main.py index 0faf3de21..9b2e77c56 100644 --- a/tests/algorithms/test_dispatch_parity_vs_main.py +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -534,7 +534,7 @@ def test_rloo_advantage_adapter_is_the_grpo_broadcast(): def _gae_inputs(): - """Fresh tensors per call: `advantage_gae` mutates `kl` in place (`k *= ...`).""" + """Fresh tensors per call: `advantage_gae` mutates `kl` in place.""" kl = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.4, 0.5])] values = [torch.tensor([0.5, 0.25, 0.125]), torch.tensor([1.0, 2.0])] return dict( @@ -595,7 +595,7 @@ def test_gae_adapter_matches_mains_numbers(cp_disabled, kl_coef, gamma, lambd): def test_gae_terminal_reward_lands_on_the_last_token(cp_disabled): - """`k[-1] += reward` is the whole reward signal; dropping it trains on KL alone.""" + """The terminal reward is the whole signal; without it, only KL trains.""" from relax.algorithms.advantages import compute_advantages_and_returns args = _args("ppo", kl_coef=0.0, gamma=1.0, lambd=1.0)