From 309cd81912bbfbe6be43a5daaaac6bb3345000d2 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 13:11:21 +0800 Subject: [PATCH 01/22] 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 5401070b0cb5f6b8d88f2a4b61e7f5f8c4070212 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 13:13:46 +0800 Subject: [PATCH 02/22] feat(gdpo): add the GDPO multi-reward advantage estimator GDPO (arXiv 2601.05242) standardises each reward component within its prompt group before combining them. A group whose rollouts share the same summed reward but differ in their components -- (correct, badly formatted) and (wrong, well formatted) both sum to 1 -- carries no signal under GRPO and is discarded; GDPO keeps both components' signal. Three steps, split across the two stages that already exist: steps 1 and 2 (per-component group standardisation, weighted combine) run in the reward normaliser on the rollout side and collapse to one scalar per sample, so the TransferQueue schema is unchanged. Step 3 (batch whitening) runs in the advantage stage, where the data-parallel group exists. Step 3 is whitened per training batch, not per merged rollout: the caller concatenates `num_rollout_minis` batches before the advantage stage, and whitening across the join centres both against a shared mean -- 4 of 8 samples flipped sign in a hand-checked case, which is a different objective rather than a rounding difference. `_whiten_by_segment` splits them back using `mini_batch_sizes`, so Eq. 6 holds regardless of `num_rollout_minis`. The six capability flags added here all have GDPO as their consumer: `forbids_normalize_advantages` (step 3 already whitens per sequence), `requires_rewards_normalization`, `supports_fully_async=False` (that mode's single-replica service sees one slice and a slice of one sample whitens to zero -- silently), `uses_reward_components`, `min_group_size=2` (an unbiased group std is undefined for one sample), and `allows_reward_post_process_hooks` (both hooks return ahead of the normaliser and would skip steps 1 and 2). Numerics: GDPO_EPS is 1e-4, matching the reference implementation, and deliberately not the 1e-6 the GRPO path is frozen at. Collapse is detected by exact equality rather than a tolerance -- any relative tolerance wide enough to catch float error also discards the genuinely informative [10000, 10000.005, 10000.010]. `distributed_mean_std` is two-pass in float64; the one-pass form returns variance 0 for [1000.0, 1000.01, 1000.02, 1000.03]. Verified on a single H100: Ray Job succeeded, 4 optimizer steps, loss / grad_norm / advantages all finite, no NaN or Inf (RELAX_REVISION=3caa34cc4c4b8c2f0c9aa654ddb911e1a639d9a2). rollout/advantages averaging 0 is by construction; the non-zero pg_loss (0.18-0.34) and grad_norm (1.29-2.22) are what show the gradient is real. Tests: 1637 passed, same 2 failures + 2 errors as main@98a1274 on this machine. --- docs/en/examples/algorithms.md | 76 ++ docs/en/guide/adding-an-algorithm.md | 13 +- docs/en/guide/configuration.md | 2 +- docs/zh/examples/algorithms.md | 76 ++ docs/zh/guide/adding-an-algorithm.md | 10 +- docs/zh/guide/configuration.md | 2 +- examples/algorithms/README.md | 24 +- examples/gdpo/README.md | 74 ++ examples/gdpo/__init__.py | 1 + examples/gdpo/reward_gdpo.py | 76 ++ examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh | 137 ++++ .../post_process_genrm_swap.py | 15 +- relax/algorithms/advantages.py | 96 +++ relax/algorithms/numerics.py | 171 +++++ relax/algorithms/policy.py | 2 +- relax/algorithms/rewards.py | 135 +++- relax/algorithms/spec.py | 70 ++ relax/backends/megatron/loss.py | 19 +- relax/components/advantages.py | 6 + relax/utils/arguments.py | 108 +++ relax/utils/types.py | 20 + relax/utils/utils.py | 12 +- tests/algorithms/test_advantage_estimators.py | 178 ++++- tests/algorithms/test_algorithm_registry.py | 16 +- tests/algorithms/test_algos_roles.py | 6 +- .../algorithms/test_arguments_spec_driven.py | 177 ++++- .../test_dispatch_parity_vs_main.py | 30 + .../algorithms/test_distributed_whitening.py | 102 +++ tests/algorithms/test_example_reward_gdpo.py | 215 ++++++ tests/algorithms/test_gdpo.py | 677 ++++++++++++++++++ tests/algorithms/test_policy_loss_dispatch.py | 2 +- .../test_post_process_rewards_dispatch.py | 4 +- 32 files changed, 2493 insertions(+), 59 deletions(-) create mode 100644 examples/gdpo/README.md create mode 100644 examples/gdpo/__init__.py create mode 100644 examples/gdpo/reward_gdpo.py create mode 100644 examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh create mode 100644 relax/algorithms/numerics.py create mode 100644 tests/algorithms/test_distributed_whitening.py create mode 100644 tests/algorithms/test_example_reward_gdpo.py create mode 100644 tests/algorithms/test_gdpo.py diff --git a/docs/en/examples/algorithms.md b/docs/en/examples/algorithms.md index 370e5adda..3584482fc 100644 --- a/docs/en/examples/algorithms.md +++ b/docs/en/examples/algorithms.md @@ -277,6 +277,81 @@ SAPO_ARGS=( --- +## GDPO + +GDPO (Group reward-Decoupled Normalization Policy Optimization, [arXiv 2601.05242](https://arxiv.org/abs/2601.05242)) targets **multi-reward** training. It standardizes each reward component within its prompt group and only then combines them, instead of summing the rewards first and normalizing once as GRPO does. + +### How It Works + +For prompt $i$ with $G$ rollouts and $n$ reward components: + +**Step 1 — per-reward group standardization:** + +$$A_k^{(i,j)} = \frac{r_k^{(i,j)} - \mathrm{mean}_j\{r_k^{(i,\cdot)}\}}{\mathrm{std}_j\{r_k^{(i,\cdot)}\} + \epsilon}$$ + +**Step 2 — weighted sum:** + +$$A_\text{sum}^{(i,j)} = \sum_k w_k A_k^{(i,j)}$$ + +The weights multiply the **normalized** advantages, not the raw rewards. After step 1 every component is on the same scale, so a weight expresses relative importance rather than the component's units. + +**Step 3 — batch-wise whitening:** + +$$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\mathrm{std}_\text{batch} + \epsilon}$$ + +**Why this beats GRPO:** when one component is constant across a group (reward collapse), GRPO's summed reward collapses too, the whole group's advantages go to zero, and the samples are wasted. Under GDPO only *that component* contributes zero while the others still carry signal. + +**On $\epsilon$:** GDPO uses $\epsilon = 10^{-4}$ at both steps, matching the reference implementation (the `scale_rewards` GDPO branch of TRL's `GRPOTrainer`), whereas GRPO / GSPO / SAPO / CISPO keep this repository's existing $10^{-6}$. The two only diverge on near-degenerate groups: with binary rewards and a group of 8 the within-group standard deviation is around 0.4 and the constants differ by 0.02%, but a continuous reward (the paper's maths setup scores response length) can leave a group at a standard deviation of ~$10^{-3}$, where $10^{-4}$ damps that group's signal by about 7% against 0.08% for $10^{-6}$. Groups that collapse *exactly* never reach this division; they are detected by exact equality and zeroed. + +### Key Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--advantage-estimator gdpo` | — | Enable GDPO | +| `--gdpo-reward-keys` | — | **Required**, at least two. Keys in the reward dict to standardize independently, e.g. `correctness format` | +| `--gdpo-reward-weights` | all 1.0 | Per-component weights; length must match `--gdpo-reward-keys` | +| `--reward-key` | — | **Required**; selects the scalar used for metrics and the `raw_reward` column | +| `--n-samples-per-prompt` | — | Must be >= 2 (the unbiased group std is undefined at $G=1$) | + +The reward function must return a dict containing every key. A missing key, a non-numeric value, a bool, or NaN/Inf raises rather than defaulting to 0.0 — a silently zeroed component is indistinguishable from a genuinely collapsed one. + +### Quick Start + +```bash +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --gdpo-reward-weights 1.0 1.0 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + --n-samples-per-prompt 8 +) +``` + +A complete runnable example lives in [`examples/gdpo/`](https://github.com/redai-infra/Relax/tree/main/examples/gdpo). + +### Known Deviations + +Two differences between this implementation and the paper. Confirm they are acceptable before training. Step 3's batch boundary used to be a third; it has since been corrected — see below. + +**Step 3's batch boundary (now aligned).** Eq. 6 normalises over one training batch. The caller merges `num_rollout_minis` of them with `concat_rollout_batches` before the advantage stage, so step 3 has to be told where the boundaries are. They travel in `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY`, which all three actor paths (colocate and hybrid) set; `loss.py` forwards them to the advantage dispatcher as `mini_batch_sizes`, and **only GDPO reads it** — every other estimator absorbs it in `**_unused` and is bit-identical either way. Each segment all-reduces across the data-parallel group, so the statistics cover both a whole training batch and every rank. Why it matters: whitening merged batches centres them all on a pooled mean, and on a measured example four of eight samples **change sign** — a different objective, not a precision difference. + +**`--fully-async` remains unsupported** and is rejected during argument validation: it hands advantage computation to the single-replica Advantages deployment, which has no data-parallel group, never sees the batch boundaries, and consumes one `global_batch_size / num_iters_per_train_update` slice at a time — when that quotient is 1 the whitened output is identically zero and the run trains on no signal at all, quietly. + +1. **A single reward does not reduce to GRPO.** Step 3 still applies, leaving a positive scalar difference from GRPO (data-dependent, measured around 1.21). Use `--advantage-estimator grpo` if you want GRPO semantics. +2. **$G=2$ discards magnitude.** Any two distinct values standardize to exactly $\pm 1/\sqrt{2}$, so with a group of two the only thing distinguishing components is their weights. + +### Mutually Exclusive Options + +- `--normalize-advantages`: step 3 already whitens per sequence; adding the token-level pass on top is not meaningful. +- `--custom-reward-post-process-path`: that hook short-circuits reward post-processing entirely, silently skipping steps 1 and 2 while the run still reports itself as GDPO. +- `--agentic-custom-advantage-path`: the second early return in `post_process_rewards`, which likewise returns ahead of the normalizer, with the same consequence. One flag, `AlgorithmSpec.allows_reward_post_process_hooks`, guards both. +- `--fully-async`: see above. + +All of these fail during argument validation. Combining it with `--dynamic-sampling-filter-path` logs a warning instead: the built-in `check_reward_nonzero_std` judges a group by the single `--reward-key` scalar and may drop groups whose signal lives in the other components. + +--- + ## Algorithm Comparison | Algorithm | Advantage Computation | Policy Loss | KL Constraint | @@ -289,6 +364,7 @@ SAPO_ARGS=( | **GSPO** | Group-relative reward | PPO-Clip + sequence-level KL | Sequence-level ratio | | **SAPO** | Group-relative reward | Sigmoid gate | Temperature-controlled | | **RLOO** | Leave-one-out baseline | Unclipped REINFORCE | Optional KL loss (same as GRPO) | +| **GDPO** | Per-reward group standardization + weighted sum + batch whitening | PPO-Clip (hard clip) | Optional KL loss | ## Next Steps diff --git a/docs/en/guide/adding-an-algorithm.md b/docs/en/guide/adding-an-algorithm.md index b415cf13d..199984fb5 100644 --- a/docs/en/guide/adding-an-algorithm.md +++ b/docs/en/guide/adding-an-algorithm.md @@ -38,11 +38,11 @@ Honestly: **not "one dict entry".** |---|---| | 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) | +| Also needs new command-line options (as GDPO needs `--gdpo-reward-keys`) | 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. +scattered if/elif chains — not the cost of adding an algorithm. GDPO is in the +last row. The `ALGOS` role table is the one part that genuinely costs nothing: it derives itself from the registry. @@ -81,6 +81,9 @@ Capability fields: | `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 | +| `supports_fully_async` | Set `False` to reject `--fully-async`, where advantages are computed slice-by-slice in a single-replica service with no data-parallel group | +| `allows_reward_post_process_hooks` | Set `False` to block both `--custom-reward-post-process-path` and `--agentic-custom-advantage-path`; each returns from `post_process_rewards` ahead of the normalizer and would silently skip your reward stage | +| `uses_reward_components` | The algorithm consumes several named reward components rather than one scalar; drives the `--gdpo-reward-keys` validation | The four `validate_*` functions in `relax/utils/arguments.py` consume every field in that table except `kl_level`, `needs_full_log_probs` and @@ -110,7 +113,7 @@ 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 +TransferQueue schema fixed — multi-reward algorithms such as GDPO collapse their components to a scalar here. **Advantage estimator** (`relax/algorithms/advantages.py`), signature @@ -175,3 +178,5 @@ is exactly what this registry exists to remove. ## References - [Algorithm Reference](../examples/algorithms.md) +- GDPO is the most recent algorithm to go through this process; read + `relax/algorithms/` alongside `examples/gdpo/`. diff --git a/docs/en/guide/configuration.md b/docs/en/guide/configuration.md index 0f8e54028..0e9d12545 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 | 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 | +| `--advantage-estimator` | str | grpo | generated from `ALGORITHM_SPECS` in `relax/algorithms/spec.py`; currently `grpo`, `gspo`, `sapo`, `cispo`, `rloo`, `gdpo`, `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/examples/algorithms.md b/docs/zh/examples/algorithms.md index 76a8b8231..2e41a5260 100644 --- a/docs/zh/examples/algorithms.md +++ b/docs/zh/examples/algorithms.md @@ -274,6 +274,81 @@ SAPO_ARGS=( --- +## GDPO + +GDPO（Group reward-Decoupled Normalization Policy Optimization，[arXiv 2601.05242](https://arxiv.org/abs/2601.05242)）面向**多奖励**训练。它对每个奖励分量分别做组内标准化，再合并——而不是像 GRPO 那样先把多个奖励加起来再归一化。 + +### 算法原理 + +设第 $i$ 个 prompt 采样 $G$ 条 rollout，共 $n$ 个奖励分量。 + +**第一步 —— 逐奖励组内标准化**： + +$$A_k^{(i,j)} = \frac{r_k^{(i,j)} - \mathrm{mean}_j\{r_k^{(i,\cdot)}\}}{\mathrm{std}_j\{r_k^{(i,\cdot)}\} + \epsilon}$$ + +**第二步 —— 加权求和**： + +$$A_\text{sum}^{(i,j)} = \sum_k w_k A_k^{(i,j)}$$ + +注意权重乘在**归一化后的 advantage** 上，不是乘在原始 reward 上。经过第一步各分量已在同一尺度，权重表达的是相对重要性，而不是分量的量纲。 + +**第三步 —— batch 级白化**： + +$$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\mathrm{std}_\text{batch} + \epsilon}$$ + +**相对 GRPO 的收益**：当一组 rollout 的各**分量不同、但总和恰好相同**时（例如 `(1,0)` 与 `(0,1)` 都求和为 1），GRPO 看到的组内总奖励无差异、整组 advantage 归零被丢弃；GDPO 对每个分量单独标准化，仍能保留各分量的学习信号。若某个分量在组内恒定，则只有**该分量**贡献 0、其它分量照常提供信号；若**所有**分量都恒定，GDPO 与 GRPO 一样返回零。 + +**关于 $\epsilon$**：GDPO 的两步都用 $\epsilon = 10^{-4}$，与参考实现（TRL `GRPOTrainer` 的 `scale_rewards` GDPO 分支）一致，而 GRPO / GSPO / SAPO / CISPO 沿用本仓库既有的 $10^{-6}$。两者的差别只在近乎塌缩的组上显现：二值 reward、组大小 8 时组内标准差约 0.4，两个取值的差异是 0.02%；但连续 reward（论文的数学实验用响应长度）可能让某组的标准差落到 $10^{-3}$ 量级，此时 $10^{-4}$ 会把该组的信号额外压低约 7%，而 $10^{-6}$ 只压低 0.08%。**完全**塌缩的组不会走到这个除法——它们由 exact 相等判定后直接置零。 + +### 关键参数 + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `--advantage-estimator gdpo` | — | 启用 GDPO | +| `--gdpo-reward-keys` | — | **必填**，至少两个。奖励函数返回的 dict 中要独立归一化的 key，如 `correctness format` | +| `--gdpo-reward-weights` | 全 1.0 | 各分量权重，长度须与 `--gdpo-reward-keys` 一致 | +| `--reward-key` | — | **必填**，选出用于 metrics 与 `raw_reward` 列的标量 | +| `--n-samples-per-prompt` | — | 必须 ≥ 2（组内无偏标准差在 $G=1$ 时无定义） | + +奖励函数必须返回包含全部 key 的 dict。缺 key、非数值、bool、NaN/Inf 都会直接报错而不是填 0——静默填 0 会把契约违约伪装成真实的 reward collapse。 + +### 快速开始 + +```bash +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --gdpo-reward-weights 1.0 1.0 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + --n-samples-per-prompt 8 +) +``` + +完整可运行示例见 [`examples/gdpo/`](https://github.com/redai-infra/Relax/tree/main/examples/gdpo)。 + +### 已知偏差 + +以下两点是实现与论文之间的实际差异，训练前请确认可以接受。第三步的 batch 边界曾经也在此列，现已修正——见下。 + +**第三步的 batch 边界（已对齐）**。论文 Eq. 6 在**一个训练批**上归一化。调用方会先把 `num_rollout_minis` 个训练批用 `concat_rollout_batches` 合并再进 advantage 阶段，所以第三步必须知道批边界。边界由 `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY` 携带（colocate 与 hybrid 三条路径都写入），`loss.py` 作为 `mini_batch_sizes` 传给 advantage 分发器，**只有 GDPO 消费**——其余估计器由 `**_unused` 吞掉，逐位不变。每段各自跨 DP all-reduce，所以统计量既覆盖完整训练批、也覆盖全部 rank。这一点为什么重要：合并白化会把两个批都对着共同均值中心化，实测 8 个样本里有 4 个**符号翻转**——那是另一个优化目标，不是精度差异。 + +**`--fully-async` 仍不受支持**，参数校验阶段直接拒绝：那条路径把 advantage 计算交给单副本的 Advantages 服务，它没有数据并行通信域，也拿不到批边界，且每次只消费 `global_batch_size / num_iters_per_train_update` 的一个切片；当这个商为 1 时白化输出恒为 0，训练会安静地在零信号上跑完。 + +1. **单个奖励时 GDPO 不退化为 GRPO**。第三步仍然生效，结果与 GRPO 相差一个正标量（与数据相关，实测约 1.21）。要 GRPO 语义就直接用 `--advantage-estimator grpo`。 +2. **$G=2$ 时幅度信息丢失**。任意两个不同值经无偏标准化后恒为 $\pm 1/\sqrt{2}$，此时分量之间的区分度只来自权重。 + +### 互斥项 + +- 不能与 `--normalize-advantages` 同用：第三步已经做过序列级白化，再叠加 token 级白化没有意义。 +- 不能与 `--custom-reward-post-process-path` 同用：该钩子会整段短路奖励后处理，导致第一、二步被静默跳过，而训练日志仍然显示算法是 GDPO。 +- 不能与 `--agentic-custom-advantage-path` 同用：`post_process_rewards` 里的第二个早返回点，同样赶在归一化器之前返回，后果与上一条相同。两者由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守。 +- 不能与 `--fully-async` 同用（见上）。 + +以上都会在参数校验阶段直接报错。配合 `--dynamic-sampling-filter-path` 时会给出警告：内置的 `check_reward_nonzero_std` 只看 `--reward-key` 那一个标量，可能丢掉只存在于其它分量的信号。 + +--- + ## 算法对比 | 算法 | Advantage 计算 | 策略损失 | KL 约束方式 | @@ -286,6 +361,7 @@ SAPO_ARGS=( | **GSPO** | 组相对奖励 | PPO-Clip + 序列级 KL | 序列级 ratio | | **SAPO** | 组相对奖励 | Sigmoid 门控 | 温度控制 | | **RLOO** | Leave-one-out 基线 | 非裁剪 REINFORCE | 可选 KL loss（同 GRPO） | +| **GDPO** | 逐奖励组内标准化 + 加权求和 + batch 白化 | PPO-Clip（硬裁剪） | 可选 KL loss | ## 下一步 diff --git a/docs/zh/guide/adding-an-algorithm.md b/docs/zh/guide/adding-an-algorithm.md index ebeb5f9d8..72b671944 100644 --- a/docs/zh/guide/adding-an-algorithm.md +++ b/docs/zh/guide/adding-an-algorithm.md @@ -26,9 +26,9 @@ relax/algorithms/ |---|---| | 复用现成的 reward 归一化 / advantage / policy loss，只是组合方式不同 | 1 个（`spec.py`） | | 需要一种新的数学（如新的 advantage 公式） | 2–3 个（`spec.py` + 对应的实现模块） | -| 还需要新的命令行参数 | 4–6 个（上述 + `arguments.py` 的参数声明与校验 + 示例 + 文档） | +| 还需要新的命令行参数（如 GDPO 的 `--gdpo-reward-keys`） | 4–6 个（上述 + `arguments.py` 的参数声明与校验 + 示例 + 文档） | -注册表消除的是「同一个算法名散落在 6 处 if/elif」，不是「新增算法零成本」。 +注册表消除的是「同一个算法名散落在 6 处 if/elif」，不是「新增算法零成本」。GDPO 走的是最后一档。 `ALGOS` 角色表是唯一真正做到零改动的部分——它从注册表自动派生。 @@ -64,6 +64,9 @@ relax/algorithms/ | `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`。适用于没有重要性比值修正的目标函数 | +| `supports_fully_async` | 设为 `False` 可拒绝 `--fully-async`（该模式下 advantage 由单副本服务按切片计算，无 DP 通信域） | +| `allows_reward_post_process_hooks` | 设为 `False` 可同时拦住 `--custom-reward-post-process-path` 与 `--agentic-custom-advantage-path`——这两个钩子都会在归一化器之前从 `post_process_rewards` 返回，静默跳过本算法的奖励阶段 | +| `uses_reward_components` | 算法消费多个具名奖励分量而非单个标量，驱动 `--gdpo-reward-keys` 校验 | 表里除 `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` 里读的：新增一个前所未有的取值需要在那里加分支，复用已有取值则不用。 @@ -82,7 +85,7 @@ def normalize_my_strategy(args, samples, raw_rewards): REWARD_NORMALIZERS["my_strategy"] = normalize_my_strategy ``` -产出必须是**每个 sample 一个标量**。这条约束让 TransferQueue 的 schema 保持不变——即使算法内部要看多个奖励分量，也要在这一层收敛成一个标量。 +产出必须是**每个 sample 一个标量**。这条约束让 TransferQueue 的 schema 保持不变——多奖励算法（如 GDPO）也是在这一层把各分量收敛成一个标量的。 **Advantage 估计器**（`relax/algorithms/advantages.py`），签名 `fn(args, *, rewards, kl, loss_masks, response_lengths, total_lengths, values) -> (advantages, returns)`，两者都是 `list[Tensor]`： @@ -125,3 +128,4 @@ pytest tests/algorithms/ -v ## 参考 - [算法参考](../examples/algorithms.md) +- GDPO 是最近一个走完整个流程的例子，可以对照 `relax/algorithms/` 与 `examples/gdpo/` 阅读。 diff --git a/docs/zh/guide/configuration.md b/docs/zh/guide/configuration.md index 957029903..9059aaa4d 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 | 由 `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 系数启用 | +| `--advantage-estimator` | str | grpo | 由 `relax/algorithms/spec.py` 的 `ALGORITHM_SPECS` 生成，当前为 `grpo`、`gspo`、`sapo`、`cispo`、`rloo`、`gdpo`、`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/algorithms/README.md b/examples/algorithms/README.md index 3ee0f407b..27287938d 100644 --- a/examples/algorithms/README.md +++ b/examples/algorithms/README.md @@ -18,6 +18,7 @@ Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator | **CISPO** | `--advantage-estimator cispo` | 保留梯度方向、需要更高精度 | | **GSPO** | `--advantage-estimator gspo` | 序列级约束、稳定训练 | | **SAPO** | `--advantage-estimator sapo` | 平滑优化、soft 信任域 | +| **GDPO** | `--advantage-estimator gdpo` | 多奖励、分量独立归一化 | ## 选择建议 @@ -68,6 +69,13 @@ Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator - 梯度流更平滑，避免梯度突变 - 适合对稳定性要求高的场景 +### GDPO（多奖励） + +- 每个 reward 分量分别做组内标准化，再加权合并 +- 某个分量组内塌缩时，其它分量仍保留学习信号（GRPO 会丢掉整组） +- 适合 correctness + format、correctness + length 这类多目标任务 +- 详见 [examples/gdpo/](../gdpo/README.md) + ## 快速开始 ### 基础操作：修改算法参数 @@ -158,7 +166,7 @@ bash scripts/training/text/run-qwen3-4B-8xgpu.sh | 参数 | 默认值 | 说明 | | ----------------------- | -------------------- | --------------------------------------------------------------------------------------------------------------- | -| `--advantage-estimator` | `grpo` | 算法类型：`grpo`, `cispo`, `gspo`, `sapo`, `ppo`, `rloo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline` | +| `--advantage-estimator` | `grpo` | 算法类型：`grpo`, `cispo`, `gspo`, `sapo`, `rloo`, `gdpo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`（取值由 `relax/algorithms/spec.py` 的注册表生成，以 `--help` 为准） | | `--eps-clip` | `0.2` | 下方裁剪边距（ratio 下界 = `1 - eps_clip`） | | `--eps-clip-high` | 与 `--eps-clip` 相同 | 上方裁剪边距（ratio 上界 = `1 + eps_clip_high`） | | `--clip-grad` | — | 梯度裁剪范数，CISPO 下推荐设为 `1.0` | @@ -194,6 +202,17 @@ bash scripts/training/text/run-qwen3-4B-8xgpu.sh | `--sapo-tau-pos` | `1.0` | Positive advantage 的温度参数 | | `--sapo-tau-neg` | `1.05` | Negative advantage 的温度参数（更高 = 更强抑制） | +### GDPO 专用参数 + +| 参数 | 默认值 | 说明 | +| ------------------------ | -------- | --------------------------------------------------------------- | +| `--gdpo-reward-keys` | — | **必填**，至少两个。奖励 dict 中要独立归一化的 key | +| `--gdpo-reward-weights` | 全 `1.0` | 各分量权重，长度须与 keys 一致；乘在**归一化后**的 advantage 上 | +| `--reward-key` | — | **必填**，选出 metrics 与 `raw_reward` 用的标量 | +| `--n-samples-per-prompt` | — | 必须 ≥ 2 | + +GDPO 与 `--normalize-advantages`、`--custom-reward-post-process-path`、`--agentic-custom-advantage-path` 和 `--fully-async` 互斥，参数校验阶段会报错。 + ### PPO 专用参数 | 参数 | 默认值 | 说明 | @@ -259,6 +278,7 @@ GSPO_ARGS=( examples/algorithms/ ├── README.md (本文件) ├── run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh (CISPO 多模态示例) +├── ../gdpo/ (GDPO 双奖励示例) ├── ... (其他算法脚本) ``` @@ -272,6 +292,7 @@ examples/algorithms/ - **CISPO**：需要精细学习信号时更好，但需要 KL 约束 - **GSPO**：长序列任务，训练更稳定 - **PPO**：如果已有 Critic 资源，性能可能更好 +- **GDPO**：多个奖励分量各自需要归一化时用它 ### Q: CISPO 的梯度波动很大，正常吗？ @@ -293,3 +314,4 @@ examples/algorithms/ - [REINFORCE++ - Simple Efficient Alignment](https://arxiv.org/abs/2501.03262) - [RLOO - Back to Basics (Ahmadian et al. 2024)](https://arxiv.org/abs/2402.14740) - [RLOO - Buy 4 REINFORCE Samples, Get a Baseline for Free (Kool et al. 2019)](https://arxiv.org/abs/1905.12705) +- [GDPO - Group reward-Decoupled Normalization](https://arxiv.org/abs/2601.05242) diff --git a/examples/gdpo/README.md b/examples/gdpo/README.md new file mode 100644 index 000000000..7216e3a48 --- /dev/null +++ b/examples/gdpo/README.md @@ -0,0 +1,74 @@ +# GDPO 示例：correctness + format 双奖励 + +本目录是 GDPO（[arXiv 2601.05242](https://arxiv.org/abs/2601.05242)）的最小可运行示例，用 Qwen3-0.6B 单卡在 GSM8K 上训练。 + +## 为什么需要 GDPO + +奖励函数 `reward_gdpo.py` 返回两个分量： + +- `correctness` —— `` 标签里的答案是否正确（0 或 1） +- `format` —— 输出是否同时带 `` 与 ``（0、0.5 或 1） + +这两个分量会**不同步**：模型可能答对但没按格式输出，也可能格式完美但答错。 + +GRPO 把它们**加起来**再做一次组内归一化。问题是不同的分量组合可能得到**相同的总和**：一组 rollout 里，有的是「答对但格式差」`(correctness=1, format=0)`、有的是「答错但格式好」`(0, 1)`，总奖励都等于 1——GRPO 看到组内总奖励全相同，整组 advantage 归零、样本白采，可两个分量各自明明都有信号。 + +GDPO 分别对每个分量做组内标准化再合并，就能区分这两类样本，`correctness` 与 `format` 各自的差异都会转成梯度信号。（反过来，若两个分量在组内**都**恒定，GDPO 与 GRPO 一样返回零，不会无中生有。） + +## 运行 + +```bash +export MODEL_DIR=/path/to/models # 需含 Qwen3-0.6B +export DATA_DIR=/path/to/data # 需含 gsm8k/train.jsonl +export EXP_DIR=/path/to/experiments + +bash examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh +``` + +### 数据要求 + +每条 prompt 需要**自带**格式要求，否则基座模型不会产出 ``/`` 标签，`format` 分量会恒为 0（组内塌缩），GDPO 就退化成只看 `correctness`。准备数据时给 question 追加一句即可： + +```python +instruction = ( + "\n\nThink step by step inside tags, then give only the " + "final number inside tags." +) +df["question"] = df["question"] + instruction +``` + +**不要用 `--system-prompt` 代替**：`relax/utils/data/data_utils.py:181` 把 system message 的 content 构造成多模态 list（`content: [{"type": "text", ...}]`），Qwen3-0.6B 这类纯文本 chat template 渲染时会报 +`TypeError: can only concatenate str (not "list") to str`。这是既有的框架限制，与 GDPO 无关。 + +## 参数说明 + +```bash +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format # 独立归一化的分量，至少两个 + --gdpo-reward-weights 1.0 1.0 # 可省略，默认全 1 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score # 必填：metrics 与 raw_reward 用的标量 + --n-samples-per-prompt 8 # 必须 ≥ 2 +) +``` + +`--gdpo-reward-weights` 乘的是**归一化之后**的 advantage，不是原始 reward。经过第一步各分量已经是单位方差，所以权重表达的是相对重要性，与分量本身的量纲无关——把 `format` 的取值范围从 `[0,1]` 改成 `[0,100]` 不会改变训练结果。 + +## 换成自己的奖励 + +改 `reward_gdpo.py` 的 `compute_gdpo_reward`，返回的 dict 需要包含 `--gdpo-reward-keys` 列出的全部 key，外加 `--reward-key` 指定的那个标量。 + +分量缺失、非数值、bool、NaN/Inf 都会直接报错。这是有意的：静默填 0 会让一个坏掉的奖励函数看起来像是「这一维恰好塌缩了」，训练照跑，问题要很久以后才暴露。 + +## 已知偏差 + +1. **第三步的 batch 边界（已正确处理）**。调用方为效率会先合并多个训练批再调用 advantage，但 `_whiten_by_segment` 用 `mini_batch_sizes` 把它们切回**每个 optimizer 训练批各自白化**，因此 `num_rollout_minis > 1` 时仍对齐论文 Eq. 6，**不要求** `rollout_batch_size × n_samples_per_prompt == global_batch_size`。本脚本把 `4 × 8` 与 `--global-batch-size 32` 设成相等只是让例子最简单，并非必需。跨 DP 的 all-reduce 保证统计量覆盖全部 rank。**`--fully-async` 会在参数校验阶段被拒绝**——那条路径的切片可能小到只有一个样本，白化输出恒为 0。 +2. **单个奖励时 GDPO 不等于 GRPO**。step1 除以 `std_g + 1e-4`、GRPO 除以 `std_g + 1e-6`，各组 `std_g` 不同 → 尺度因子逐组不同，step3 还会再做一次 batch 白化，所以不是「差一个正标量」那么简单。要 GRPO 语义就用 `--advantage-estimator grpo`。 +3. **`--n-samples-per-prompt 2` 时幅度信息丢失**：任意两个不同值标准化后恒为 ±0.7071。示例用 8 就是为了避开这一点。 + +## 冲突项 + +`--normalize-advantages`、`--custom-reward-post-process-path`、`--agentic-custom-advantage-path` 和 `--fully-async` 都不能与 GDPO 同用，参数校验阶段会直接报错。第一个会造成双重白化；第二、三个都会赶在归一化器之前从 `post_process_rewards` 返回，导致 GDPO 的前两步被静默跳过（这两个由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守）；第四个的统计窗口可能小到只剩一个样本。 + +配 `--dynamic-sampling-filter-path` 时只警告不报错：内置过滤器按 `--reward-key` 的单个标量判组，可能丢掉只在其它分量里有信号的组。 diff --git a/examples/gdpo/__init__.py b/examples/gdpo/__init__.py new file mode 100644 index 000000000..9f3863608 --- /dev/null +++ b/examples/gdpo/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. diff --git a/examples/gdpo/reward_gdpo.py b/examples/gdpo/reward_gdpo.py new file mode 100644 index 000000000..476d6cee0 --- /dev/null +++ b/examples/gdpo/reward_gdpo.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Two-component reward for the GDPO example: correctness and format. + +GDPO standardizes each component within its prompt group before combining them, +so a group whose rollouts differ in their reward *components* but share the same +*summed* reward still carries a learning signal. Example: (correct, badly +formatted) and (wrong, well formatted) both sum to 1 -- GRPO sees one constant +summed reward and the whole group contributes nothing, while GDPO keeps each +component's signal. (If every component is constant within the group, GDPO +returns zero too.) + +Wire it up with:: + + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + +``score`` is what ``--reward-key`` selects for metrics and for the ``raw_reward`` +column; the training signal comes from the two components. +""" + +import re +from typing import Any + + +_ANSWER_RE = re.compile(r"(.*?)", re.DOTALL) +_THINK_RE = re.compile(r".*?", re.DOTALL) + + +def _extract_answer(response: str) -> str | None: + match = _ANSWER_RE.search(response) + return match.group(1).strip() if match else None + + +def _final_answer(label: Any) -> str: + """Normalise a GSM8K label to just the final answer. + + GSM8K ships ``answer`` as the full worked solution ending in ``#### 36``, so + comparing a model's ``36`` against the whole string makes + ``correctness`` zero for every rollout. That collapses the component in every + group, and GDPO silently degrades to the single ``format`` reward -- which is + the one thing this example exists to demonstrate it does not do. + + Handling it here rather than only in a data-prep step keeps the example + working against the dataset it names. Labels without the marker pass through + unchanged, so a pre-cleaned dataset behaves identically. + """ + text = str(label).strip() + return text.rsplit("####", 1)[-1].strip() if "####" in text else text + + +def compute_gdpo_reward(response: str, label: Any) -> dict[str, float]: + """Score one response on answer correctness and on output format. + + The two components are deliberately decorrelated: a response can be correct + without the expected tags, and well-formatted while wrong. That is the + situation GDPO handles better than a summed reward. + """ + answer = _extract_answer(response) + correctness = 1.0 if answer is not None and answer == _final_answer(label) else 0.0 + + has_think = _THINK_RE.search(response) is not None + format_score = 0.5 * float(has_think) + 0.5 * float(answer is not None) + + return { + "score": correctness, + "correctness": correctness, + "format": format_score, + } + + +async def reward_func(args: Any, sample: Any, **kwargs: Any) -> dict[str, float]: + """Entry point for ``--custom-rm-path``.""" + return compute_gdpo_reward(sample.response, sample.label) diff --git a/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh new file mode 100644 index 000000000..447d2f555 --- /dev/null +++ b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh @@ -0,0 +1,137 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3-0.6B single-GPU GDPO training on GSM8K. +# +# GDPO (arXiv 2601.05242) standardizes each reward component within its prompt +# group before combining them, so a group whose rollouts share the same summed +# reward but differ in their components still carries signal -- e.g. (correct, +# badly formatted) and (wrong, well formatted) both sum to 1, which GRPO flattens +# to nothing. The reward function in reward_gdpo.py returns both components; +# --gdpo-reward-keys names them. +# +# Usage: +# bash examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh + +set -ex +set -o pipefail + +now=$(date "+%Y-%m-%d-%H:%M:%S") + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# Auto-source local environment when not launched via an external entrypoint +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../scripts/entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" + +PROJECT_NAME="${PROJECT_NAME:=relax-gdpo}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" +NUM_ROLLOUT="${NUM_ROLLOUT:=20}" + +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3-0.6B + --ref-load ${MODEL_DIR}/Qwen3-0.6B + --load ${EXP_DIR}/Qwen3-0.6B_mcore_gdpo/ + --save ${EXP_DIR}/Qwen3-0.6B_mcore_gdpo/ + --save-interval 100 + --max-actor-ckpt-to-keep 1 + --megatron-to-hf-mode bridge +) + +# NOTE: the format instruction belongs in the prompt text, not in +# --system-prompt. relax/utils/data/data_utils.py:181 builds the system message +# with multimodal list content (`content: [{"type": "text", ...}]`), which a +# text-only chat template such as Qwen3-0.6B's cannot render: +# TypeError: can only concatenate str (not "list") to str +# Prepare the dataset so each question already asks for the / +# tags; see examples/gdpo/README.md. +ROLLOUT_ARGS=( + --prompt-data ${DATA_DIR}/gsm8k/train.jsonl + --input-key question + --label-key answer + --apply-chat-template + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 4 + --n-samples-per-prompt 8 + --rollout-max-response-len 1024 + --rollout-temperature 1.0 + # 4 * 8 == 32 just keeps this example minimal. GDPO's step 3 stays aligned + # with Eq. 6 regardless: _whiten_by_segment splits any merged rollout back + # into per-optimizer-batch segments and whitens each on its own, so + # num_rollout_minis > 1 does NOT require rollout_batch_size * n == gbs. + --global-batch-size 32 +) + +# The two reward components come from examples/gdpo/reward_gdpo.py. +# --reward-key selects the scalar used for metrics and the raw_reward column; +# --gdpo-reward-keys names the components GDPO standardizes independently. +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --gdpo-reward-weights 1.0 1.0 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + --eps-clip 0.2 + --kl-coef 0.00 + --entropy-coef 0.00 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --clip-grad 1.0 +) + +WANDB_ARGS=( + --use-tensorboard + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name qwen3-0.6b-gdpo-gpu1-${now} +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.5 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +mkdir -p log +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 1], "rollout": [1, 1]}' \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --use-health-check \ + --balance-data \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GDPO_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen3-0.6b-gdpo-gpu1-${now}.log diff --git a/examples/generate_reward_model/post_process_genrm_swap.py b/examples/generate_reward_model/post_process_genrm_swap.py index f3556b99d..8a61290f1 100644 --- a/examples/generate_reward_model/post_process_genrm_swap.py +++ b/examples/generate_reward_model/post_process_genrm_swap.py @@ -109,12 +109,15 @@ def _grpo_normalize(args, raw_rewards): 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. + rather than asking ``spec.is_group_normalized``. That property is also true + of ``gdpo_decoupled``, which standardises several reward components + independently and is not what the code below computes. GDPO cannot reach + here today (``allows_reward_post_process_hooks=False`` rejects this hook in + ``relax/utils/arguments.py``), so this is a guard against a future + multi-reward algorithm that allows the hook, not a live bug. 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 + this reimplementation needs a human to look at it. """ from relax.algorithms import get_algorithm diff --git a/relax/algorithms/advantages.py b/relax/algorithms/advantages.py index a93129ad6..781a6dc08 100644 --- a/relax/algorithms/advantages.py +++ b/relax/algorithms/advantages.py @@ -17,7 +17,9 @@ from typing import Any, Callable import torch +import torch.distributed as dist +from relax.algorithms.numerics import GDPO_EPS, distributed_mean_std, is_collapsed from relax.algorithms.spec import get_algorithm from relax.utils.training.ppo_utils import ( get_advantages_and_returns_batch, @@ -27,6 +29,42 @@ ) +def whiten_scalar(values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None) -> torch.Tensor: + """Sequence-level whitening of one scalar per sample. + + This is GDPO's batch-wise normalisation (arXiv 2601.05242, Eq. 6). It is + deliberately *not* the token-level ``distributed_masked_whiten`` used by + ``--normalize-advantages``: weighting by token count would let long + responses dominate the statistics, which Eq. 6 does not do. + + ``process_group`` must be supplied wherever the caller holds only a shard of + the values. In the Megatron path each data-parallel rank owns + ``num_rollout_minis * global_batch_size / dp_size`` samples — its shard of + the whole rollout, merged before ``compute_advantages_and_returns`` runs — so + whitening locally would give every rank its own mean and scale, not the "one + global scale factor" the maths assumes. ``None`` means "I own every value you + need"; note the single-replica ``Advantages`` deployment cannot be that caller + for GDPO, since ``supports_fully_async=False`` rejects the configuration that + would route here — so in practice ``None`` is the CPU-side and test callers. + + Scope of one call: this whitens exactly the tensor it is handed. When the + caller has merged several training batches, :func:`_whiten_by_segment` splits + them back and calls this once per optimizer batch, so each call stays aligned + with Eq. 6's per-batch statistic regardless of ``num_rollout_minis``. Only + the un-segmented path (``mini_batch_sizes=None``) whitens a whole merged + tensor at once. + + A batch where every value is identical returns exact zeros; see + :func:`relax.algorithms.numerics.is_collapsed`. + """ + if is_collapsed(values, process_group=process_group): + return torch.zeros_like(values) + mean, std = distributed_mean_std(values, process_group=process_group) + if not torch.isfinite(std): + return torch.zeros_like(values) + return (values - mean) / (std + GDPO_EPS) + + 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) @@ -41,6 +79,51 @@ def advantage_grpo_broadcast(args: Any, *, rewards, kl, **_unused): return advantages, returns +def _whiten_by_segment(values, mini_batch_sizes, process_group): + """Whiten each training batch separately, in the order they were merged. + + Every rank runs the same number of segments -- ``num_rollout_minis`` comes + from the minibatch plan rather than from the data -- so the per-segment + collectives stay matched across the data-parallel group. + """ + if mini_batch_sizes is None: + return whiten_scalar(values, process_group=process_group) + # Validate whatever was passed, including a single segment: treating an empty + # or malformed list as "fall back to one window" would silently restore the + # merged behaviour this function exists to replace. + if not mini_batch_sizes or any(not isinstance(n, int) or n <= 0 for n in mini_batch_sizes): + raise ValueError(f"mini_batch_sizes must be a non-empty list of positive ints, got {mini_batch_sizes}.") + if sum(mini_batch_sizes) != values.numel(): + raise ValueError( + f"mini_batch_sizes {mini_batch_sizes} sum to {sum(mini_batch_sizes)}, " + f"but this rank holds {values.numel()} samples." + ) + if len(mini_batch_sizes) == 1: + return whiten_scalar(values, process_group=process_group) + out, start = [], 0 + for size in mini_batch_sizes: + out.append(whiten_scalar(values[start : start + size], process_group=process_group)) + start += size + return torch.cat(out) + + +def advantage_gdpo(args: Any, *, rewards, kl, process_group=None, mini_batch_sizes=None, **_unused): + """GDPO step 3: whiten the combined per-sample advantage, then broadcast. + + Steps 1 and 2 (per-reward group standardisation and the weighted sum) ran + on the rollout side, so ``rewards`` already holds one ``A_sum`` per sample. + Doing step 3 here rather than there keeps it out of reach of ``--custom- + reward-post-process-path``, which short-circuits reward post-processing + entirely, and off the streaming transfer-batch boundary with its undersized + tail flushes. + """ + reward_tensor = _as_reward_tensor(rewards, kl) + whitened = _whiten_by_segment(reward_tensor, mini_batch_sizes, process_group) + returns = get_grpo_returns(whitened, kl) + advantages = list(returns) + 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).""" @@ -123,6 +206,7 @@ def advantage_gae( ADVANTAGE_FNS: dict[str, Callable[..., tuple[list[torch.Tensor], list[torch.Tensor]]]] = { "grpo_broadcast": advantage_grpo_broadcast, + "gdpo": advantage_gdpo, "reinforce_plus_plus": advantage_reinforce_plus_plus, "reinforce_plus_plus_baseline": advantage_reinforce_plus_plus_baseline, "gae": advantage_gae, @@ -139,9 +223,19 @@ def compute_advantages_and_returns( total_lengths: list[int] | None = None, values: list[torch.Tensor] | None = None, padded_total_lengths: list[int] | None = None, + process_group: dist.ProcessGroup | None = None, + mini_batch_sizes: list[int] | None = None, ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: """Dispatch to the estimator registered for ``args.advantage_estimator``. + ``mini_batch_sizes`` is this rank's per-training-batch sample counts, in the + order the caller merged them. Only estimators whose statistics are defined + per batch read it; the rest absorb it in ``**_unused`` and are unaffected. + + ``process_group`` is the group across which the batch is sharded, or + ``None`` when the caller holds every sample. Estimators that compute batch- + level statistics need it to see the whole batch. + ``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 @@ -158,4 +252,6 @@ def compute_advantages_and_returns( total_lengths=total_lengths, values=values, padded_total_lengths=padded_total_lengths, + process_group=process_group, + mini_batch_sizes=mini_batch_sizes, ) diff --git a/relax/algorithms/numerics.py b/relax/algorithms/numerics.py new file mode 100644 index 000000000..8a8676aae --- /dev/null +++ b/relax/algorithms/numerics.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Shared numerical constants and guards for the algorithm implementations. + +Both the reward stage and the advantage stage standardise values by dividing by +a standard deviation, so they need the same epsilon and the same notion of +"this group carries no signal". Keeping those here prevents the two stages +from drifting apart. +""" + +import torch +import torch.distributed as dist + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +_LOGGED_GROUP_SIZE = False + + +def _log_group_once(process_group: dist.ProcessGroup | None) -> None: + """Report the reduction group once per process. + + Whether a batch statistic is global or per-shard is invisible in the loss + curve: a misconfigured run where tensor parallelism ate the extra GPUs + leaves the data-parallel group at size 1, the all-reduce becomes an + identity, and everything still trains. This line is what makes that + distinguishable in a log. + """ + global _LOGGED_GROUP_SIZE + if _LOGGED_GROUP_SIZE: + return + _LOGGED_GROUP_SIZE = True + if process_group is None: + logger.info("Batch statistics are local (no process group); the caller owns the whole batch.") + else: + logger.info( + "Batch statistics reduce over dp_world=%d (this rank is dp_rank=%d).", + dist.get_world_size(process_group), + dist.get_rank(process_group), + ) + + +STD_EPS = 1e-6 +"""Epsilon added to a 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 the whole point of the equivalence tests, so this +constant is not free to move. +""" + +GDPO_EPS = 1e-4 +"""Epsilon for GDPO's two standardisation steps. + +Deliberately not :data:`STD_EPS`. The reference implementation +(``trl/trainer/grpo_trainer.py``, the ``scale_rewards`` GDPO path) divides by +``std + 1e-4`` at both the per-reward group step and the batch step, and GDPO +is new here, so there is no prior Relax behaviour that matching it would break. + +The choice only bites near-degenerate groups: with binary rewards and eight +samples the group std is around 0.4 and the two constants differ by 0.02%, but +a continuous reward (the paper's maths setup uses response length) can leave a +group with std ~1e-3, where 1e-6 and 1e-4 disagree by roughly 10% on the scale +factor. Exactly-collapsed groups never reach either constant; they are caught +by :func:`is_collapsed` and zeroed. +""" + + +def is_collapsed(values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None) -> bool: + """Whether every value is identical, i.e. the spread carries no signal. + + Tested by exact equality rather than by comparing the standard deviation + against a tolerance. A tolerance has to be relative to the magnitude (the + mean of N equal float32 values does not come back exactly equal to them, so + a collapsed group still shows std ~= 1e-8 times its magnitude), and any + relative tolerance large enough to catch that also erases real signal: with + ``std <= 1e-6 * max|x|``, the perfectly informative batch + ``[10000, 10000.005, 10000.010, 10000.015]`` is thrown away. Exact equality + has no such false positives, and near-equality is already damped by the + caller's epsilon (:data:`GDPO_EPS` on the GDPO whitening path, + :data:`STD_EPS` on the group path). + """ + if process_group is None: + if values.numel() == 0: + return True + return bool(values.min() == values.max()) + + # Every rank in the group has to reach the collective, including one whose + # shard came out empty — returning early there would hang the others. An + # empty shard contributes -inf to both halves, which is the identity for MAX + # and therefore leaves the reduction to the ranks that do have samples. + empty = values.numel() == 0 + neg_infinity = torch.tensor(float("-inf"), dtype=values.dtype, device=values.device) + bounds = torch.stack( + [neg_infinity if empty else -values.min(), neg_infinity if empty else values.max()], + ) + dist.all_reduce(bounds, op=dist.ReduceOp.MAX, group=process_group) + low, high = -bounds[0], bounds[1] + if not torch.isfinite(low): + return True # every rank was empty + return bool(low == high) + + +def collapsed_columns(values: torch.Tensor, dim: int) -> torch.Tensor: + """Per-column version of :func:`is_collapsed` for a ``[G, K]`` group. + + Returns a boolean tensor of shape ``[K]``: ``True`` where that reward + component took the same value across the whole group. + """ + return values.amax(dim=dim) == values.amin(dim=dim) + + +def distributed_mean_std( + values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + """Mean and unbiased std of ``values``, optionally across + ``process_group``. + + Each rank holds its own shard of the batch, so a local ``std()`` would give + every rank a different scale factor. Reducing across the group makes the + statistics describe the whole batch, which is what the framework already + does for ``--normalize-advantages`` (see + ``relax.utils.distributed_utils.distributed_masked_whiten``). + + Two passes, in float64. The one-pass form ``E[x^2] - E[x]^2`` subtracts two + nearly equal large numbers when the values sit far from zero, and the result + is dominated by rounding: for ``[1000.0, 1000.01, 1000.02, 1000.03]`` it + returns a variance of exactly 0 (true std 1.29e-2), and for the tighter + ``[10000.0, 10000.001, 10000.002, 10000.003]`` it returns std 4.6 instead of + 1.3e-3 — off by a factor of 3660. Scaling the first example up to 1e4 + without tightening it just reproduces the zero, not the inflation. + Neither is loud. The first silently zeroes every advantage in the batch; + the second silently rescales them. Reward magnitudes like these are + ordinary: the GDPO paper's own maths setup uses a length reward, and token + counts live in the thousands. + + The extra collective is two scalars, which is not worth optimising away. + """ + _log_group_once(process_group) + + # float64 throughout: the cancellation above is a precision problem, and + # doing the arithmetic in the caller's float32 reintroduces it even with the + # two-pass formula. + work = values.double() + count = torch.tensor(float(work.numel()), dtype=torch.float64, device=work.device) + total = work.sum() + + if process_group is not None: + first = torch.stack([count, total]) + dist.all_reduce(first, op=dist.ReduceOp.SUM, group=process_group) + count, total = first[0], first[1] + + if count == 0: + zero = torch.zeros((), dtype=values.dtype, device=values.device) + return zero, zero + + mean = total / count + # Centred before squaring, so no large offset survives into the sum. + sum_sq_dev = (work - mean).pow(2).sum() + if process_group is not None: + dist.all_reduce(sum_sq_dev, op=dist.ReduceOp.SUM, group=process_group) + + # Bessel-corrected, matching torch.std()'s default so the single-reward GDPO + # scale factor stays derivable from the GRPO group statistics. + variance = sum_sq_dev / torch.clamp(count - 1, min=1.0) + # variance cannot be negative now that it is a sum of squares; the clamp is + # only guarding the exactly-zero case against a -0.0. + std = torch.sqrt(torch.clamp(variance, min=0.0)) + return mean.to(values.dtype), std.to(values.dtype) diff --git a/relax/algorithms/policy.py b/relax/algorithms/policy.py index cc033c9fa..12e023bbc 100644 --- a/relax/algorithms/policy.py +++ b/relax/algorithms/policy.py @@ -23,7 +23,7 @@ def policy_loss_ppo_clip(args: Any, *, log_probs, ppo_kl, advantages): - """Standard clipped surrogate objective (GRPO, GSPO, PPO, REINFORCE++).""" + """Standard clipped surrogate objective (GRPO, GSPO, GDPO, REINFORCE++).""" return compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index 9cc7efd22..e1c28a6fe 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -6,24 +6,24 @@ ``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. +schema — even for multi-reward algorithms, which collapse their components to a +single scalar here. """ +import math +from numbers import Real from typing import Any, Callable import torch +from relax.algorithms.numerics import GDPO_EPS, STD_EPS, collapsed_columns +from relax.utils.logging_utils import get_logger 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. +logger = get_logger(__name__) -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. -""" +GROUP_EPS = STD_EPS def group_positions(samples: list[Any], expected_size: int) -> dict[int, list[int]]: @@ -129,10 +129,129 @@ def normalize_group_leave_one_out(args: Any, samples: list[Any], raw_rewards: li return normalized_rewards.tolist() +def resolve_gdpo_keys(args: Any) -> list[str]: + """The reward components GDPO normalises independently.""" + keys = list(getattr(args, "gdpo_reward_keys", None) or []) + if len(keys) < 2: + raise ValueError(f"--gdpo-reward-keys needs at least two reward keys, got {keys}.") + duplicates = {key for key in keys if keys.count(key) > 1} + if duplicates: + raise ValueError(f"--gdpo-reward-keys contains duplicates: {sorted(duplicates)}.") + return keys + + +def resolve_gdpo_weights(args: Any, keys: list[str]) -> list[float]: + """Per-component weights, defaulting to 1.0 each. + + The weights multiply the *normalised* advantages (arXiv 2601.05242, Eq. 7), + not the raw rewards. That is the point: after step 1 every component is on + the same scale, so a weight expresses relative importance rather than + accidentally encoding the component's units. + """ + weights = getattr(args, "gdpo_reward_weights", None) + if weights is None: + return [1.0] * len(keys) + if len(weights) != len(keys): + raise ValueError(f"--gdpo-reward-weights has {len(weights)} entries but --gdpo-reward-keys has {len(keys)}.") + resolved = [float(w) for w in weights] + for key, weight in zip(keys, resolved, strict=True): + # argparse happily parses "nan" and "inf" for a float option. Unchecked, + # the weighted sum turns non-finite, whiten_scalar reads a non-finite std + # as a collapse, and the batch silently produces zero advantages. + if not math.isfinite(weight): + raise ValueError(f"--gdpo-reward-weights for {key!r} is {weight}, which is not finite.") + if all(weight == 0.0 for weight in resolved): + # Every component gets multiplied by zero, so the combined advantage is + # identically zero: the run trains on no signal and still exits cleanly. + # The same shape of failure as the non-finite case above, one line later. + raise ValueError(f"--gdpo-reward-weights are all zero ({resolved}); the combined advantage would be 0.") + return resolved + + +def extract_reward_components(samples: list[Any], keys: list[str]) -> torch.Tensor: + """Build the ``[B, K]`` component matrix, rejecting malformed rewards. + + Contract violations raise instead of defaulting to 0.0. A silently zeroed + component is indistinguishable from a genuinely collapsed one, so falling + back would hide a broken reward function behind plausible-looking training. + """ + rows: list[list[float]] = [] + for position, sample in enumerate(samples): + values = sample.get_reward_components(keys) + row: list[float] = [] + for key, value in zip(keys, values, strict=True): + # `numbers.Real` rather than `(int, float)`: reward functions routinely + # return numpy scalars, and only np.float64 happens to subclass float — + # np.float32/np.int64 would be rejected while np.float64 sailed through. + # bool and np.bool_ stay rejected (bool subclasses int; np.bool_ is not + # a Real), because a boolean reward is almost always a mistake. + if isinstance(value, bool) or not isinstance(value, Real): + raise TypeError( + f"Reward {key!r} of sample {position} must be a real number, " + f"got {value!r} ({type(value).__name__})." + ) + numeric = float(value) + if not math.isfinite(numeric): + raise ValueError(f"Reward {key!r} of sample {position} is {numeric}, which is not finite.") + row.append(numeric) + rows.append(row) + return torch.tensor(rows, dtype=torch.float32) + + +def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: + """GDPO steps 1 and 2 (arXiv 2601.05242, Eq. 4 and Eq. 7). + + Step 1 standardises each reward component within its prompt group; step 2 + combines them with the configured weights. The result is one scalar per + sample, so it travels through the existing ``rewards`` column and needs no + TransferQueue schema change. Step 3 (batch whitening) runs later, in + :func:`relax.algorithms.advantages.advantage_gdpo`. + + Standardising per component before combining is what separates GDPO from + GRPO: when one component collapses within a group, only that component + contributes zero, while GRPO's summed reward would collapse and discard the + whole group. + """ + keys = resolve_gdpo_keys(args) + weights = resolve_gdpo_weights(args, keys) + + components = extract_reward_components(samples, keys) + positions_by_group = group_positions(samples, args.n_samples_per_prompt) + + normalized = torch.zeros_like(components) + fully_collapsed_groups = 0 + for positions in positions_by_group.values(): + group = components[positions] + centered = group - group.mean(dim=0, keepdim=True) + std = group.std(dim=0) + collapsed = collapsed_columns(group, dim=0) + scaled = centered / (std + GDPO_EPS) + normalized[positions] = torch.where(collapsed.unsqueeze(0), torch.zeros_like(scaled), scaled) + if bool(collapsed.all()): + fully_collapsed_groups += 1 + + if fully_collapsed_groups == len(positions_by_group): + # Every component collapsed in every group, so this batch produces no + # gradient at all. Usually the reward function is constant for these + # prompts (e.g. a format reward when nothing in the prompt asks for a + # format). Worth one line, because the symptom downstream is simply + # "loss does not move". + logger.warning( + "GDPO: all reward components collapsed in all %d groups of this batch (keys=%s); " + "the batch contributes no gradient. Check that each of these rewards actually varies " + "across rollouts of the same prompt.", + fully_collapsed_groups, + keys, + ) + + weight_tensor = torch.tensor(weights, dtype=torch.float32) + return (normalized * weight_tensor).sum(dim=1).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, + "gdpo_decoupled": normalize_gdpo_decoupled, } diff --git a/relax/algorithms/spec.py b/relax/algorithms/spec.py index d124845db..7c76d9285 100644 --- a/relax/algorithms/spec.py +++ b/relax/algorithms/spec.py @@ -121,6 +121,57 @@ class AlgorithmSpec: only sees a slice of). Those get their own field; do not fold them here. """ + supports_fully_async: bool = True + """Whether the algorithm is correct under ``--fully-async``. + + That mode routes advantage computation to the single-replica + ``relax.components.advantages`` deployment, which owns no data-parallel + group and consumes one ``global_batch_size / num_iters_per_train_update`` + slice at a time. An algorithm whose advantages depend on batch-level + statistics computes them over that slice instead of the batch, and at + slice size 1 gets no signal at all — silently, since the run still + converges and exits cleanly. Note ``--hybrid`` is *not* affected: it uses + the colocate role set, so advantages are computed in the Megatron worker + where the data-parallel group exists. + + Do not be tempted to relax this into "allow it when the slice happens to + equal the batch". The slice the deployment actually receives also depends + on which TransferQueue sampler the controller installed + (``relax/core/controller.py``), and that choice is global to every + consumer. Under ``--balance-data`` it is ``SeqlenBalancedSampler``, whose + ``batch_size`` is *per data-parallel rank*; the deployment passes no + ``sampling_config``, so it would receive one rank's share of a + token-balanced split rather than the batch. + """ + + uses_reward_components: bool = False + """Whether the algorithm consumes several named reward components rather + than the single scalar ``--reward-key`` selects. Drives the + ``--gdpo-reward-keys`` / ``--gdpo-reward-weights`` validation. + + Those two options stay GDPO-prefixed on purpose: every algorithm-specific + option in Relax names its algorithm (``--sapo-tau-pos``, + ``--disable-grpo-std-normalization``), and GDPO is so far the only member + of this category — renaming now would mean guessing the abstraction from a + single example. When a second ``uses_reward_components`` algorithm lands, + rename both to algorithm-neutral options and keep the old spellings as + deprecated aliases for one release, the way ``--loss-type sft_loss`` is + handled in ``relax/utils/arguments.py``. + """ + + min_group_size: int = 1 + allows_reward_post_process_hooks: bool = True + """False when a user-supplied reward hook would silently disable the + algorithm's own reward stage. + + Guards both short-circuits in ``relax.utils.utils.post_process_rewards``: + ``--custom-reward-post-process-path`` replaces the function outright, and + ``--agentic-custom-advantage-path`` returns before the normaliser runs. + They are one flag rather than two because an algorithm that cannot tolerate + one cannot tolerate the other -- ``reinforce_plus_plus_baseline`` already + rejects both by hand in ``relax/utils/arguments.py``, and GDPO needs the + same pair.""" + @property def is_group_normalized(self) -> bool: """Whether rewards get normalised per prompt group on the rollout @@ -175,6 +226,25 @@ def is_group_normalized(self) -> bool: requires_global_token_loss=True, requires_on_policy_updates=True, ), + "gdpo": AlgorithmSpec( + name="gdpo", + reward_normalizer="gdpo_decoupled", + advantage_fn="gdpo", + policy_loss_fn="ppo_clip", + # Step 3 already whitens per sequence; --normalize-advantages would add a + # second, token-level pass on top of it. + forbids_normalize_advantages=True, + requires_rewards_normalization=True, + # Step 3 needs the training batch. Under --fully-async it would only ever + # see one slice of it, and a slice of one sample yields zero advantages. + supports_fully_async=False, + uses_reward_components=True, + # Step 1 divides by an unbiased group std, undefined for a single sample. + min_group_size=2, + # Either reward hook short-circuits reward post-processing, which would + # silently skip steps 1 and 2 while the run still reports itself as GDPO. + allows_reward_post_process_hooks=False, + ), "ppo": AlgorithmSpec( name="ppo", reward_normalizer="none", diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index bce3e4f98..fc7aafe64 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -12,6 +12,7 @@ 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.backends.megatron.data import ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY 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 ( @@ -523,8 +524,9 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) `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, + listed here, because the list this replaced had already gone stale (it never + gained "gdpo") and 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. @@ -587,6 +589,19 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) # at the padded offsets; omitting it does not raise, it reads the wrong # token positions. padded_total_lengths=padded_total_lengths, + # rollout_data is this rank's shard of the WHOLE rollout, not of one + # training batch: actor.py collects `num_rollout_minis` windows of + # global_batch_size/dp_size and concat_rollout_batches merges them before + # this call ("we may need normalize the whole rollout", actor.py). So the + # reduction below makes the statistic describe the rollout across the DP + # group -- which spans every optimizer step in it, not one batch. See the + # known-deviation note in docs/*/examples/algorithms.md. + process_group=mpu.get_data_parallel_group(), + # Per-training-batch counts, written by actor.py before it merges the + # rollout. Passing them is what lets a batch-level statistic describe one + # batch instead of the whole merged rollout; estimators that do not have + # one absorb this in **_unused. + mini_batch_sizes=rollout_data.get(ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY), ) # Optional pure OPD mode: remove all non-OPD reward contribution. diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 4cfb69b75..bb49c1748 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -178,6 +178,12 @@ def compute_advantages_and_returns(self, rollout_data: Dict[str, Any]) -> Dict[s # `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. + # + # This deployment is a single replica and owns every sample it was + # handed, so batch statistics need no reduction. Note the batch is + # one `global_batch_size / num_iters_per_train_update` slice, not + # the whole training batch — see the GDPO notes in the docs. + process_group=None, ) # Optional pure OPD mode: remove all non-OPD reward contribution. diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 64f9dc33d..b1d4fe54f 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -1749,6 +1749,28 @@ def add_algo_arguments(parser): default=1.05, help="Temperature for negative advantages in SAPO (default: 1.05)", ) + parser.add_argument( + "--gdpo-reward-keys", + type=str, + nargs="+", + default=None, + help=( + "Names of the reward components GDPO standardizes independently, e.g. " + "`--gdpo-reward-keys correctness format`. The reward function must return a dict " + "containing every key. At least two keys are required." + ), + ) + parser.add_argument( + "--gdpo-reward-weights", + type=float, + nargs="+", + default=None, + help=( + "Per-component weights for GDPO, matching --gdpo-reward-keys in length. " + "Defaults to 1.0 each. The weights multiply the *normalized* advantages, " + "not the raw rewards, so they express relative importance rather than units." + ), + ) parser.add_argument( "--disable-compute-advantages-and-returns", action="store_false", @@ -3005,6 +3027,92 @@ def validate_algorithm_args(args) -> None: f"The {spec.name!r} advantage estimator requires advantage normalization. " "Please add `--normalize-advantages` to your command." ) + if spec.forbids_normalize_advantages and args.normalize_advantages: + raise ValueError( + f"The {spec.name!r} advantage estimator already whitens advantages per sequence; " + "`--normalize-advantages` would apply a second, token-level whitening on top. " + "Please remove it." + ) + if spec.requires_rewards_normalization and not args.rewards_normalization: + raise ValueError( + f"The {spec.name!r} advantage estimator needs reward normalization. " + "Please remove `--disable-rewards-normalization`." + ) + if not spec.allows_reward_post_process_hooks: + # Both hooks return from post_process_rewards before the registry's + # normalizer runs, so either one silently skips this algorithm's reward + # stage while the run still reports itself as that algorithm. + for option, value in ( + ("--custom-reward-post-process-path", args.custom_reward_post_process_path), + ("--agentic-custom-advantage-path", getattr(args, "agentic_custom_advantage_path", None)), + ): + if value is not None: + raise ValueError( + f"`{option}` short-circuits reward post-processing, which would silently skip " + f"{spec.name!r}'s reward normalization while the run still reports itself as " + f"{spec.name!r}. Please drop one of the two." + ) + if args.n_samples_per_prompt < spec.min_group_size: + raise ValueError( + f"The {spec.name!r} advantage estimator needs `--n-samples-per-prompt` >= " + f"{spec.min_group_size}, got {args.n_samples_per_prompt}." + ) + # `--hybrid` also ends up with fully_async set, but only later in validation and + # only as an implementation detail: it uses the colocate role set, so advantages + # are computed in the Megatron worker rather than the Advantages deployment. + # Reading both raw flags here is what distinguishes the two. + if not spec.supports_fully_async and args.fully_async and not args.hybrid: + raise ValueError( + f"The {spec.name!r} advantage estimator is not supported under --fully-async. " + "Its advantages depend on batch-level statistics, but that mode computes them in a " + "single-replica service that sees one `global_batch_size / num_iters_per_train_update` " + "slice at a time — with a slice of one sample the advantages come out all zero and the " + "run trains on no signal without failing. Use --colocate or --hybrid." + ) + + if spec.uses_reward_components: + _validate_multi_reward_args(args, spec) + + +def _validate_multi_reward_args(args, spec) -> None: + """Check the reward-component configuration for multi-reward algorithms.""" + keys = args.gdpo_reward_keys or [] + if len(keys) < 2: + raise ValueError( + f"The {spec.name!r} advantage estimator needs at least two reward keys; " + f"pass e.g. `--gdpo-reward-keys correctness format`. Got {keys}." + ) + duplicates = sorted({key for key in keys if keys.count(key) > 1}) + if duplicates: + raise ValueError(f"`--gdpo-reward-keys` contains duplicates: {duplicates}.") + + weights = args.gdpo_reward_weights + if weights is not None and len(weights) != len(keys): + raise ValueError( + f"`--gdpo-reward-weights` has {len(weights)} entries but `--gdpo-reward-keys` has {len(keys)}." + ) + + # Components arrive as a dict; without --reward-key the raw_reward column + # would hold dicts, which the TransferQueue conversion cannot represent. + if not args.reward_key: + raise ValueError( + f"The {spec.name!r} advantage estimator needs `--reward-key` to select the scalar reward " + "used for metrics and for the raw_reward column." + ) + + if args.dynamic_sampling_filter_path: + # A warning rather than an error: the filter is opt-in and a custom one may + # well be component-aware. The built-in check_reward_nonzero_std is not — it + # reads the single --reward-key scalar, so a group where that component is + # flat but another still varies gets dropped, which is exactly the case this + # estimator exists to keep. + logger.warning( + "%r combines multiple reward components, but --dynamic-sampling-filter-path filters on the " + "single --reward-key scalar (%r). Groups carrying signal only in the other components may be " + "dropped before training sees them.", + spec.name, + args.reward_key, + ) if args.n_samples_per_prompt < spec.min_group_size: raise ValueError( diff --git a/relax/utils/types.py b/relax/utils/types.py index 9c7cabb5b..4236636fe 100644 --- a/relax/utils/types.py +++ b/relax/utils/types.py @@ -170,6 +170,26 @@ def from_dict(data: dict): def get_reward_value(self, args) -> float: return self.reward if not args.reward_key else self.reward[args.reward_key] + def get_reward_components(self, keys: list[str]) -> list[Any]: + """Return the named reward components, in ``keys`` order. + + Multi-reward algorithms such as GDPO need the individual components, + which ``get_reward_value`` collapses to a single scalar. This is a + separate accessor rather than a change to that one because dynamic + sampling filters and rollout metrics depend on its current signature. + """ + if not isinstance(self.reward, dict): + raise ValueError( + f"Sample.reward must be a dict to read components {keys}, " + f"got {type(self.reward).__name__}. Make the reward function return a dict of named rewards." + ) + values = [] + for key in keys: + if key not in self.reward: + raise ValueError(f"Reward key {key!r} missing from sample reward (available: {sorted(self.reward)}).") + values.append(self.reward[key]) + return values + @property def effective_response_length(self): return sum(self.loss_mask) if self.loss_mask is not None else self.response_length diff --git a/relax/utils/utils.py b/relax/utils/utils.py index 902bc42c4..2394d87d6 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -136,7 +136,7 @@ def convert_samples_to_train_data(args: Any, samples: list[Sample] | list[list[S 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 + # --reward-key, and every GDPO run), 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"] = [ @@ -186,10 +186,12 @@ 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. + # Second short-circuit: this one replaces the normalizer wholesale, so an + # algorithm whose reward stage is load-bearing (GDPO's steps 1 and 2) would + # be silently skipped here while the run still reports itself as that + # algorithm. Argument validation rejects the combination up front, driven + # by `AlgorithmSpec.allows_reward_post_process_hooks`, which guards this + # hook and --custom-reward-post-process-path together. if getattr(args, "agentic_custom_advantage_path", None) is not None: return raw_rewards, [sample.custom_advantage for sample in samples] diff --git a/tests/algorithms/test_advantage_estimators.py b/tests/algorithms/test_advantage_estimators.py index 08a42f822..0b07feace 100644 --- a/tests/algorithms/test_advantage_estimators.py +++ b/tests/algorithms/test_advantage_estimators.py @@ -9,7 +9,11 @@ torch = pytest.importorskip("torch") -from relax.algorithms.advantages import ADVANTAGE_FNS, compute_advantages_and_returns # noqa: E402 +from relax.algorithms.advantages import ( # noqa: E402 + ADVANTAGE_FNS, + compute_advantages_and_returns, + whiten_scalar, +) def _args(estimator, **overrides): @@ -105,3 +109,175 @@ def test_reinforce_plus_plus_baseline_zeroes_masked_tokens(): 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])) + + +# ---------------- gdpo step 3 ---------------- + + +def test_whiten_scalar_produces_zero_mean_unit_std(): + out = whiten_scalar(torch.tensor([1.0, 2.0, 3.0, 4.0])) + assert abs(out.mean().item()) < 1e-6 + assert abs(out.std().item() - 1.0) < 1e-3 + + +def test_whiten_scalar_zero_variance_returns_zeros_not_noise(): + """A collapsed batch must give exactly 0, not the fp32 residual over + eps.""" + assert torch.equal(whiten_scalar(torch.full((7,), 0.7)), torch.zeros(7)) + + +def test_whiten_scalar_single_element_returns_zero(): + """std of one element is NaN under Bessel correction.""" + assert torch.equal(whiten_scalar(torch.tensor([3.0])), torch.zeros(1)) + + +def test_whiten_scalar_empty_returns_empty(): + assert whiten_scalar(torch.tensor([])).numel() == 0 + + +def test_whiten_scalar_preserves_ordering(): + values = torch.tensor([-3.0, 0.5, 0.0, 9.0]) + out = whiten_scalar(values) + assert torch.equal(values.argsort(), out.argsort()) + + +def test_gdpo_advantage_whitens_before_broadcasting(): + rewards = [2.0, -2.0] + adv, _ = compute_advantages_and_returns(_args("gdpo"), rewards=rewards, **_inputs()) + expected = whiten_scalar(torch.tensor(rewards, dtype=torch.float32)) + assert torch.allclose(adv[0], expected[0].expand(3)) + assert torch.allclose(adv[1], expected[1].expand(2)) + + +def test_gdpo_differs_from_grpo_by_a_positive_scalar(): + rewards = [1.0, -1.0, 0.5, -0.5] + inputs = _inputs(lengths=(1, 1, 1, 1)) + grpo, _ = compute_advantages_and_returns(_args("grpo"), rewards=rewards, **inputs) + gdpo, _ = compute_advantages_and_returns(_args("gdpo"), rewards=rewards, **inputs) + + grpo_flat = torch.cat(grpo) + gdpo_flat = torch.cat(gdpo) + ratio = gdpo_flat / grpo_flat + assert (ratio > 0).all() + assert torch.allclose(ratio, ratio[0].expand_as(ratio), atol=1e-4) + assert not torch.allclose(gdpo_flat, grpo_flat, atol=1e-3) + + +def test_gdpo_collapsed_batch_gives_zero_advantages(): + adv, _ = compute_advantages_and_returns(_args("gdpo"), rewards=[0.7, 0.7], **_inputs()) + assert torch.equal(adv[0], torch.zeros(3)) + assert torch.equal(adv[1], torch.zeros(2)) + + +# ---------------- collapse detection is exact ---------------- + + +def test_collapse_uses_exact_equality_not_a_relative_tolerance(): + """A relative tolerance erased this perfectly informative batch.""" + values = torch.tensor([10000.0, 10000.005, 10000.010, 10000.015]) + out = whiten_scalar(values) + assert not torch.equal(out, torch.zeros_like(out)) + assert torch.equal(values.argsort(), out.argsort()) + + +@pytest.mark.parametrize("magnitude", [0.0, 0.1, 0.7, 1.0, 1e6, -3.5]) +def test_identical_values_collapse_exactly_at_any_magnitude(magnitude): + out = whiten_scalar(torch.full((7,), magnitude)) + assert torch.equal(out, torch.zeros(7)) + + +def test_two_values_differing_by_one_ulp_are_not_collapsed(): + base = torch.tensor(1.0) + values = torch.stack([base, torch.nextafter(base, torch.tensor(2.0))] * 2) + out = whiten_scalar(values) + # The name's claim is the > 0: is_collapsed tests exact equality, so a + # one-ULP spread is real signal and must survive. An earlier version of this + # test asserted only the upper bound, which an all-zero output also passes -- + # i.e. it did not test the thing it was named after. + assert out.abs().max() > 0.0 + # eps damps it towards zero rather than amplifying it to unit scale. + assert out.abs().max() < 1.0 + + +def test_empty_batch_is_treated_as_collapsed(): + assert whiten_scalar(torch.tensor([])).numel() == 0 + + +# ---------------- distributed statistics ---------------- + + +def test_distributed_mean_std_without_a_group_matches_torch(): + from relax.algorithms.numerics import distributed_mean_std + + values = torch.tensor([1.0, 2.0, 4.0, 8.0]) + mean, std = distributed_mean_std(values) + assert torch.allclose(mean, values.mean()) + assert torch.allclose(std, values.std(), atol=1e-5) + + +def test_whiten_scalar_matches_manual_formula_without_a_group(): + values = torch.tensor([1.0, 2.0, 4.0, 8.0]) + # 1e-4, not the GRPO path's 1e-6: whiten_scalar is GDPO step 3 and matches + # the reference implementation's epsilon. Hardcoded rather than imported so + # the test would notice the constant moving. + expected = (values - values.mean()) / (values.std() + 1e-4) + assert torch.allclose(whiten_scalar(values), expected, atol=1e-6) + + +def test_sharded_whitening_differs_from_local_whitening(): + """Why the process group matters: local stats give each shard its own + scale. + + Simulates two DP ranks by whitening each shard alone and comparing against + whitening the concatenated batch. + """ + shard_a = torch.tensor([-0.7, 0.7]) + shard_b = torch.tensor([-1.4, 1.4]) + + local = torch.cat([whiten_scalar(shard_a), whiten_scalar(shard_b)]) + joint = whiten_scalar(torch.cat([shard_a, shard_b])) + + # Local whitening flattens the two shards onto the same amplitude; the joint + # statistics keep shard_b's larger relative contribution. + assert torch.allclose(local[:2].abs(), local[2:].abs(), atol=1e-4) + assert not torch.allclose(joint[:2].abs(), joint[2:].abs(), atol=1e-2) + + +def test_compute_advantages_and_returns_accepts_a_process_group_kwarg(): + """Both call sites pass it; every estimator must tolerate it.""" + from relax.algorithms import list_algorithm_names + from relax.algorithms.spec import get_algorithm + + # "gae" needs a critic and "reinforce_plus_plus" imports megatron.core for the + # CP world size; neither is available on a CPU-only runner. + needs_megatron = {"gae", "reinforce_plus_plus"} + for name in list_algorithm_names(): + if get_algorithm(name).advantage_fn in needs_megatron: + continue + args = _args(name, kl_coef=0.0) + adv, _ = compute_advantages_and_returns(args, rewards=[0.5, -0.5], process_group=None, **_inputs()) + assert len(adv) == 2 + + +@pytest.mark.parametrize("bad", [float("inf"), float("-inf"), float("nan")]) +def test_a_non_finite_value_gives_zeros_rather_than_poisoning_the_gradient(bad): + """The `not torch.isfinite(std)` guard in whiten_scalar, which coverage + showed no test reached. + + A single inf or nan among the rewards makes the batch std non-finite. + Without the guard the division propagates nan into every advantage in the + batch and from there into the gradient, where it is far harder to trace + back. Reward functions are user code, so this is reachable input, not a + hypothetical. + """ + values = torch.tensor([1.0, bad, 2.0, 3.0]) + out = whiten_scalar(values) + assert torch.equal(out, torch.zeros_like(values)) + + +def test_a_huge_but_finite_spread_is_still_whitened(): + """The guard must catch non-finite, not merely large -- otherwise it would + silently discard batches it should scale.""" + out = whiten_scalar(torch.tensor([1e38, -1e38, 0.0])) + assert not torch.equal(out, torch.zeros_like(out)) + assert torch.isfinite(out).all() diff --git a/tests/algorithms/test_algorithm_registry.py b/tests/algorithms/test_algorithm_registry.py index 5e8706385..b3bc79b72 100644 --- a/tests/algorithms/test_algorithm_registry.py +++ b/tests/algorithms/test_algorithm_registry.py @@ -15,6 +15,7 @@ "gspo", "sapo", "cispo", + "gdpo", "ppo", "reinforce_plus_plus", "reinforce_plus_plus_baseline", @@ -62,13 +63,11 @@ def test_reward_normalizer_ids_match_current_behavior(): def test_is_group_normalized_matches_the_legacy_whitelist(): - """The pre-registry whitelist, exactly. - - group. - """ + """The pre-registry whitelist, plus GDPO which also normalizes per + 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 + assert actual == legacy | {"gdpo"} def test_gspo_is_the_only_sequence_level_kl(): @@ -104,9 +103,12 @@ def test_policy_loss_ids_match_current_behavior(): 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.min_group_size == 1 + assert spec.allows_reward_post_process_hooks is True + assert spec.requires_rewards_normalization is False + assert spec.forbids_normalize_advantages is False + assert spec.uses_reward_components is False assert spec.needs_critic is False - assert spec.requires_normalize_advantages is False def test_spec_module_has_no_heavy_imports(): diff --git a/tests/algorithms/test_algos_roles.py b/tests/algorithms/test_algos_roles.py index d5320a24b..e21f2ef3c 100644 --- a/tests/algorithms/test_algos_roles.py +++ b/tests/algorithms/test_algos_roles.py @@ -87,9 +87,9 @@ def test_every_registered_algorithm_has_a_role_mapping(): @requires_megatron -@pytest.mark.parametrize("name", ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]) +@pytest.mark.parametrize("name", ["reinforce_plus_plus", "reinforce_plus_plus_baseline", "gdpo"]) def test_previously_missing_algorithms_are_covered(name): - """These two are why the derivation exists.""" + """These three are why the derivation exists.""" assert name in ALGOS @@ -115,7 +115,7 @@ def test_ppo_keeps_its_critic(): @requires_megatron def test_policy_gradient_algorithms_have_no_critic(): - for name in ("grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus"): + for name in ("grpo", "gspo", "sapo", "cispo", "gdpo", "reinforce_plus_plus"): assert ROLES.critic not in ALGOS[name], f"{name} would start a critic service it never uses" diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py index 7c5eec96d..3764ffb82 100644 --- a/tests/algorithms/test_arguments_spec_driven.py +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -48,13 +48,15 @@ def arguments_module(monkeypatch): sys.modules.pop("relax.utils.arguments", None) -def _args(estimator="grpo", **overrides): +def _args(estimator="gdpo", **overrides): base = dict( advantage_estimator=estimator, normalize_advantages=False, rewards_normalization=True, custom_reward_post_process_path=None, n_samples_per_prompt=4, + gdpo_reward_keys=["correctness", "format"], + gdpo_reward_weights=None, reward_key="score", use_critic=False, fully_async=False, @@ -91,13 +93,48 @@ def test_validation_reads_spec_fields(): for field in ( "needs_critic", "requires_normalize_advantages", + "forbids_normalize_advantages", + "requires_rewards_normalization", + "allows_reward_post_process_hooks", + "min_group_size", + "uses_reward_components", + "supports_fully_async", ): assert field in src, f"arguments.py does not consult spec.{field}" +def test_gdpo_arguments_declared(): + src = ARGS_PATH.read_text(encoding="utf-8") + assert "--gdpo-reward-keys" in src + assert "--gdpo-reward-weights" in src + + # ---------------- behaviour ---------------- +def test_parser_accepts_gdpo_and_its_options(arguments_module): + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + + args = parser.parse_args( + [ + "--advantage-estimator", + "gdpo", + "--gdpo-reward-keys", + "correctness", + "format", + "--gdpo-reward-weights", + "1.0", + "0.5", + ] + ) + + assert args.advantage_estimator == "gdpo" + assert args.gdpo_reward_keys == ["correctness", "format"] + assert args.gdpo_reward_weights == [1.0, 0.5] + + def test_parser_rejects_an_unregistered_estimator(arguments_module): arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) parser = argparse.ArgumentParser() @@ -119,14 +156,62 @@ def test_every_registered_algorithm_is_an_accepted_choice(arguments_module): assert parsed.advantage_estimator == name +def test_gdpo_rejects_custom_reward_post_process(arguments_module): + with pytest.raises(ValueError, match="custom-reward-post-process-path"): + arguments_module.validate_algorithm_args(_args(custom_reward_post_process_path="pkg.mod.fn")) + + +def test_gdpo_rejects_normalize_advantages(arguments_module): + with pytest.raises(ValueError, match="normalize-advantages"): + arguments_module.validate_algorithm_args(_args(normalize_advantages=True)) + + +def test_gdpo_rejects_group_size_below_two(arguments_module): + with pytest.raises(ValueError, match="n-samples-per-prompt"): + arguments_module.validate_algorithm_args(_args(n_samples_per_prompt=1)) + + +def test_gdpo_rejects_disabled_rewards_normalization(arguments_module): + with pytest.raises(ValueError, match="reward normalization"): + arguments_module.validate_algorithm_args(_args(rewards_normalization=False)) + + +def test_gdpo_requires_at_least_two_keys(arguments_module): + with pytest.raises(ValueError, match="at least two reward keys"): + arguments_module.validate_algorithm_args(_args(gdpo_reward_keys=["correctness"])) + + +def test_gdpo_rejects_duplicate_keys(arguments_module): + with pytest.raises(ValueError, match="duplicates"): + arguments_module.validate_algorithm_args(_args(gdpo_reward_keys=["a", "a"])) + + +def test_gdpo_rejects_weight_count_mismatch(arguments_module): + with pytest.raises(ValueError, match="gdpo-reward-weights"): + arguments_module.validate_algorithm_args(_args(gdpo_reward_weights=[1.0])) + + +def test_gdpo_requires_reward_key(arguments_module): + with pytest.raises(ValueError, match="reward-key"): + arguments_module.validate_algorithm_args(_args(reward_key=None)) + + +def test_gdpo_happy_path(arguments_module): + args = _args() + arguments_module.validate_algorithm_args(args) + assert args.use_critic is False + + 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)) + arguments_module.validate_algorithm_args( + _args(estimator, normalize_advantages=False, gdpo_reward_keys=None) + ) def test_reinforce_family_passes_with_normalize_advantages(arguments_module): - args = _args("reinforce_plus_plus", normalize_advantages=True) + args = _args("reinforce_plus_plus", normalize_advantages=True, gdpo_reward_keys=None) arguments_module.validate_algorithm_args(args) assert args.use_critic is False @@ -134,21 +219,21 @@ def test_reinforce_family_passes_with_normalize_advantages(arguments_module): 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) + args = _args("ppo", gdpo_reward_keys=None, 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) + args = _args(estimator, gdpo_reward_keys=None, 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) +def test_non_gdpo_algorithms_ignore_reward_key_and_component_options(arguments_module): + """Only multi-reward algorithms police --reward-key / --gdpo-*.""" + args = _args("grpo", gdpo_reward_keys=None, gdpo_reward_weights=None, reward_key=None) arguments_module.validate_algorithm_args(args) @@ -168,7 +253,7 @@ def _write_yaml(tmp_path, body): 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 = _args("grpo", gdpo_reward_keys=None, reward_key=None) base.loss_type = "policy_loss" base.custom_config_path = _write_yaml(tmp_path, body) for key, value in overrides.items(): @@ -186,6 +271,29 @@ def test_yaml_cannot_switch_to_an_estimator_that_needs_a_critic(arguments_module arguments_module.apply_custom_config_overrides(args) +def test_yaml_cannot_bypass_gdpo_requirements(arguments_module, tmp_path): + """Switching to GDPO from YAML must still demand its reward keys.""" + args = _overridable_args(tmp_path, "advantage_estimator: gdpo\n") + + with pytest.raises(ValueError, match="at least two reward keys"): + arguments_module.apply_custom_config_overrides(args) + + +def test_yaml_cannot_enable_conflicting_whitening_under_gdpo(arguments_module, tmp_path): + """The dangerous case: silent double whitening rather than a crash.""" + args = _overridable_args( + tmp_path, + "normalize_advantages: true\n", + advantage_estimator="gdpo", + gdpo_reward_keys=["correctness", "format"], + reward_key="score", + n_samples_per_prompt=8, + ) + + with pytest.raises(ValueError, match="normalize-advantages"): + 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 @@ -271,7 +379,7 @@ def test_yaml_without_algorithm_changes_is_accepted(arguments_module, tmp_path): def test_no_yaml_is_a_no_op(arguments_module): - args = _args("grpo", reward_key=None) + args = _args("grpo", gdpo_reward_keys=None, reward_key=None) args.loss_type = "policy_loss" args.custom_config_path = None arguments_module.apply_custom_config_overrides(args) @@ -304,13 +412,56 @@ def test_spec_with_an_unregistered_implementation_is_rejected_at_startup(argumen 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)) + arguments_module.validate_algorithm_args(_args("grpo", gdpo_reward_keys=None, reward_key=None)) + + +# ---------------- fully-async is not a supported execution mode for GDPO ---------------- -# ---------------- fully-async ---------------- +def test_gdpo_is_rejected_under_fully_async(arguments_module): + """Otherwise it trains on one slice at a time and, at slice size 1, on + nothing.""" + with pytest.raises(ValueError, match="not supported under --fully-async"): + arguments_module.validate_algorithm_args(_args(fully_async=True)) + + +def test_gdpo_is_allowed_under_hybrid(arguments_module): + """--hybrid sets fully_async later, but uses the colocate role set: advantages + are computed in the Megatron worker, where the DP group exists.""" + arguments_module.validate_algorithm_args(_args(fully_async=True, hybrid=True)) + + +def test_gdpo_is_allowed_under_colocate(arguments_module): + arguments_module.validate_algorithm_args(_args(fully_async=False)) @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) + args = _args(estimator, fully_async=True, gdpo_reward_keys=None, reward_key=None) arguments_module.validate_algorithm_args(args) + + +def test_supports_fully_async_defaults_to_true(): + from relax.algorithms import get_algorithm, list_algorithm_names + + unsupported = {n for n in list_algorithm_names() if not get_algorithm(n).supports_fully_async} + assert unsupported == {"gdpo"} + + +def test_dynamic_sampling_filter_warns_for_multi_reward(arguments_module, caplog): + """The built-in filter judges a group by the single --reward-key scalar.""" + import logging + + with caplog.at_level(logging.WARNING): + arguments_module.validate_algorithm_args(_args(dynamic_sampling_filter_path="pkg.mod.fn")) + + assert any("dynamic-sampling-filter-path" in r.message for r in caplog.records) + + +def test_no_warning_without_a_filter(arguments_module, caplog): + import logging + + with caplog.at_level(logging.WARNING): + arguments_module.validate_algorithm_args(_args()) + + assert not any("dynamic-sampling-filter-path" in r.message for r in caplog.records) diff --git a/tests/algorithms/test_dispatch_parity_vs_main.py b/tests/algorithms/test_dispatch_parity_vs_main.py index 302b7078e..d271af6d4 100644 --- a/tests/algorithms/test_dispatch_parity_vs_main.py +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -185,6 +185,13 @@ def test_every_algorithm_main_supported_is_still_registered(): assert not missing, f"{missing} were reachable on main {MAIN_SHA[:7]} and are gone now" +def test_only_gdpo_was_added(): + """Keeps the reference tables honest: any new algorithm must be listed + here.""" + added = set(list_algorithm_names()) - set(MAIN_ALGORITHMS) + assert added == {"gdpo"}, f"unexpected new algorithms {added}; update this file's tables" + + def test_reward_normalizer_identifiers_all_resolve(): for name in list_algorithm_names(): assert get_algorithm(name).reward_normalizer in REWARD_NORMALIZERS @@ -522,3 +529,26 @@ def test_rloo_advantage_adapter_is_the_grpo_broadcast(): for left, right in zip(got, want, strict=True): assert torch.equal(left, right) + + +def test_loss_py_actually_forwards_the_mini_batch_boundaries(): + """Source-level, because loss.py needs megatron to import. + + The per-batch whitening tests all call advantage_gdpo directly, so removing + the wiring in loss.py left every one of them green while GDPO silently went + back to whitening the merged rollout. This is the assertion that fails when + that happens. + """ + import pathlib + import re + + src = (pathlib.Path(__file__).resolve().parents[2] / "relax" / "backends" / "megatron" / "loss.py").read_text( + encoding="utf-8" + ) + + call = re.search(r"compute_advantages_and_returns_impl\((.*?)\n \)", src, re.DOTALL) + assert call, "compute_advantages_and_returns_impl call not found" + assert "mini_batch_sizes=rollout_data.get(ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY)" in call.group(1), ( + "loss.py must pass the per-training-batch counts; without them GDPO's step 3 " + "normalises over the whole merged rollout again" + ) diff --git a/tests/algorithms/test_distributed_whitening.py b/tests/algorithms/test_distributed_whitening.py new file mode 100644 index 000000000..7a438e32f --- /dev/null +++ b/tests/algorithms/test_distributed_whitening.py @@ -0,0 +1,102 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""GDPO step 3 must whiten over the whole batch, not each rank's shard. + +These spawn a real two-process gloo group so the collectives are exercised for +real: a shard-local implementation, a missing collective on an empty shard, or a +mismatched call order would all show up here rather than only on a GPU cluster. +""" + +import os + +import pytest + + +torch = pytest.importorskip("torch") +mp = pytest.importorskip("torch.multiprocessing") + +import torch.distributed as dist # noqa: E402 + + +# Two ranks with deliberately different spreads: shard 1 carries twice the +# amplitude of shard 0, which only survives whitening if the statistics are +# shared. +SHARDS = [[-0.7, 0.7], [-1.4, 1.4]] + + +def _run(rank, world_size, port, mode, out): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + try: + from relax.algorithms.advantages import whiten_scalar + from relax.algorithms.numerics import distributed_mean_std, is_collapsed + + group = dist.group.WORLD + if mode == "whiten": + values = torch.tensor(SHARDS[rank], dtype=torch.float32) + out[rank] = whiten_scalar(values, process_group=group).tolist() + elif mode == "empty_shard": + # Rank 1 has nothing to contribute; rank 0 must not hang on it. + values = torch.tensor([1.0, 3.0] if rank == 0 else [], dtype=torch.float32) + out[rank] = whiten_scalar(values, process_group=group).tolist() + elif mode == "stats": + values = torch.tensor(SHARDS[rank], dtype=torch.float32) + mean, std = distributed_mean_std(values, process_group=group) + out[rank] = [mean.item(), std.item()] + elif mode == "collapsed_across_ranks": + # Identical on both ranks: only a global view can tell. + values = torch.tensor([0.7, 0.7], dtype=torch.float32) + out[rank] = [is_collapsed(values, process_group=group)] + elif mode == "collapsed_only_locally": + # Each shard is constant on its own but the batch is not. + values = torch.tensor([0.7, 0.7] if rank == 0 else [1.4, 1.4], dtype=torch.float32) + out[rank] = [is_collapsed(values, process_group=group)] + else: # pragma: no cover - guard against typos in the test itself + raise AssertionError(mode) + finally: + dist.destroy_process_group() + + +def _spawn(mode, world_size=2): + manager = mp.Manager() + out = manager.dict() + port = 29500 + abs(hash(mode)) % 2000 + mp.spawn(_run, args=(world_size, port, mode, out), nprocs=world_size, join=True) + return dict(out) + + +def test_statistics_are_shared_across_ranks(): + out = _spawn("stats") + everything = torch.tensor(SHARDS[0] + SHARDS[1]) + for rank in (0, 1): + mean, std = out[rank] + assert abs(mean - everything.mean().item()) < 1e-5, rank + assert abs(std - everything.std().item()) < 1e-4, rank + + +def test_whitening_keeps_the_larger_shard_larger(): + """Shard-local whitening flattens both shards onto the same amplitude.""" + out = _spawn("whiten") + shard0 = torch.tensor(out[0]) + shard1 = torch.tensor(out[1]) + + assert shard1.abs().max() > shard0.abs().max() * 1.5 + joint = torch.cat([shard0, shard1]) + assert abs(joint.mean().item()) < 1e-5 + assert abs(joint.std().item() - 1.0) < 1e-2 + + +def test_an_empty_shard_does_not_hang_the_group(): + """The empty rank still has to reach every collective.""" + out = _spawn("empty_shard") + assert out[1] == [] + assert len(out[0]) == 2 + assert out[0][0] < 0 < out[0][1] + + +def test_collapse_is_decided_globally_not_per_shard(): + assert _spawn("collapsed_across_ranks") == {0: [True], 1: [True]} + + per_shard = _spawn("collapsed_only_locally") + assert per_shard == {0: [False], 1: [False]}, "each shard is constant, the batch is not" diff --git a/tests/algorithms/test_example_reward_gdpo.py b/tests/algorithms/test_example_reward_gdpo.py new file mode 100644 index 000000000..b51a0ee8c --- /dev/null +++ b/tests/algorithms/test_example_reward_gdpo.py @@ -0,0 +1,215 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""The shipped GDPO example reward must produce the two configured +components.""" + +import importlib.util +import pathlib + +import pytest + + +EXAMPLE_DIR = pathlib.Path(__file__).resolve().parents[2] / "examples" / "gdpo" +REWARD_PATH = EXAMPLE_DIR / "reward_gdpo.py" +SCRIPT_PATH = EXAMPLE_DIR / "run-qwen3-0.6B-1xgpu-gdpo.sh" + + +@pytest.fixture(scope="module") +def reward_module(): + spec = importlib.util.spec_from_file_location("reward_gdpo", REWARD_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_returns_all_three_keys(reward_module): + out = reward_module.compute_gdpo_reward("x42", "42") + assert set(out) == {"score", "correctness", "format"} + + +def test_correct_and_well_formatted(reward_module): + out = reward_module.compute_gdpo_reward("reasoning42", "42") + assert out["correctness"] == 1.0 + assert out["format"] == 1.0 + + +def test_wrong_answer_but_well_formatted(reward_module): + out = reward_module.compute_gdpo_reward("reasoning7", "42") + assert out["correctness"] == 0.0 + assert out["format"] == 1.0 + + +def test_correct_answer_but_malformed(reward_module): + """The case GDPO is designed for: the components disagree.""" + out = reward_module.compute_gdpo_reward("42", "42") + assert out["correctness"] == 0.0 # no tag means the answer is unparseable + assert out["format"] == 0.0 + + +def test_partially_formatted_scores_half(reward_module): + out = reward_module.compute_gdpo_reward("42", "42") + assert out["format"] == 0.5 + assert out["correctness"] == 1.0 + + +def test_thinking_without_an_answer_scores_half_format(reward_module): + out = reward_module.compute_gdpo_reward("reasoning", "42") + assert out["format"] == 0.5 + assert out["correctness"] == 0.0 + + +def test_score_mirrors_correctness(reward_module): + for response in ("42", "7", "nothing"): + out = reward_module.compute_gdpo_reward(response, "42") + assert out["score"] == out["correctness"] + + +def test_label_is_stringified_and_stripped(reward_module): + assert reward_module.compute_gdpo_reward("42", 42)["correctness"] == 1.0 + assert reward_module.compute_gdpo_reward("42", " 42 ")["correctness"] == 1.0 + + +def test_components_are_plain_floats(reward_module): + out = reward_module.compute_gdpo_reward("42", "42") + for key in ("score", "correctness", "format"): + assert isinstance(out[key], float) + assert not isinstance(out[key], bool) + + +def test_components_survive_the_gdpo_normalizer(reward_module): + """End-to-end: example rewards feed the registered normalizer without error.""" + from types import SimpleNamespace + + from relax.algorithms.rewards import normalize_gdpo_decoupled + + responses = [ + "a42", + "b7", + "42", + "nothing at all", + ] + samples = [ + SimpleNamespace( + group_index=0, + reward=reward_module.compute_gdpo_reward(r, "42"), + get_reward_components=lambda keys, r=r: [reward_module.compute_gdpo_reward(r, "42")[k] for k in keys], + ) + for r in responses + ] + args = SimpleNamespace( + n_samples_per_prompt=4, + gdpo_reward_keys=["correctness", "format"], + gdpo_reward_weights=None, + ) + + out = normalize_gdpo_decoupled(args, samples, [0.0] * 4) + assert len(out) == 4 + assert abs(sum(out)) < 1e-4 # both components are group-centred + assert max(abs(v) for v in out) > 0.1 # and there is real signal + + +# ---------------- launch script ---------------- + + +def test_launch_script_wires_the_reward_to_the_estimator(): + src = SCRIPT_PATH.read_text(encoding="utf-8") + assert "--advantage-estimator gdpo" in src + assert "--gdpo-reward-keys correctness format" in src + assert "--custom-rm-path examples.gdpo.reward_gdpo.reward_func" in src + assert "--reward-key score" in src + + +def test_launch_script_satisfies_the_gdpo_group_size_floor(): + src = SCRIPT_PATH.read_text(encoding="utf-8") + assert "--n-samples-per-prompt 8" in src + + +def test_launch_script_does_not_enable_the_conflicting_whitening(): + src = SCRIPT_PATH.read_text(encoding="utf-8") + assert "--normalize-advantages" not in src + assert "--custom-reward-post-process-path" not in src + + +# ---------------- the label form the launch script actually feeds ---------------- +# +# The script points --prompt-data at gsm8k/train.jsonl with --label-key answer, +# and GSM8K's `answer` is the whole worked solution ending in "#### 36". Every +# test above used a pre-cleaned "42", so nothing here noticed that comparing +# against the raw string makes correctness zero for every rollout -- collapsing +# the component in every group and quietly reducing the example to its single +# format reward. + +RAW_GSM8K_LABEL = "Janet sells 16 - 3 - 4 = 9 duck eggs.\n9 * $2 = $18\n#### 18" + + +def test_correct_answer_scores_against_a_raw_gsm8k_label(reward_module): + out = reward_module.compute_gdpo_reward("work18", RAW_GSM8K_LABEL) + assert out["correctness"] == 1.0, "raw GSM8K labels must not zero out correctness" + assert out["format"] == 1.0 + + +def test_a_precleaned_label_behaves_identically(reward_module): + raw = reward_module.compute_gdpo_reward("work18", RAW_GSM8K_LABEL) + clean = reward_module.compute_gdpo_reward("work18", "18") + assert raw == clean, "normalising the label must not change a dataset that is already clean" + + +def test_a_wrong_answer_is_still_wrong_against_a_raw_label(reward_module): + """Otherwise the normalisation could be making everything match.""" + out = reward_module.compute_gdpo_reward("work99", RAW_GSM8K_LABEL) + assert out["correctness"] == 0.0 + assert out["format"] == 1.0, "format is independent of correctness; that is the point of the example" + + +def test_the_two_components_can_disagree_on_a_raw_label(reward_module): + """The case GDPO exists for, on the data the script actually loads.""" + correct_unformatted = reward_module.compute_gdpo_reward("18", RAW_GSM8K_LABEL) + wrong_formatted = reward_module.compute_gdpo_reward("w99", RAW_GSM8K_LABEL) + + assert correct_unformatted["format"] < wrong_formatted["format"] + assert correct_unformatted["correctness"] == 0.0 # no tag, so unparseable + assert wrong_formatted["correctness"] == 0.0 + + +# ---------------- the docs must not contradict the script ---------------- +# +# Three times in this branch a claim was fixed in one file and left stale in a +# sibling: the batch geometry, the "covers a full training batch" wording, and +# the algorithm-name lists. Prose cannot be diffed against code, so this checks +# the one number the prose actually depends on. + + +def _script_flag(name): + import pathlib + import re + + script = ( + pathlib.Path(__file__).resolve().parents[2] / "examples" / "gdpo" / "run-qwen3-0.6B-1xgpu-gdpo.sh" + ).read_text(encoding="utf-8") + match = re.search(rf"--{name}\s+(\d+)", script) + assert match, f"--{name} not found in the launch script" + return int(match.group(1)) + + +def test_example_geometry_satisfies_the_paper_boundary(): + """rollout_batch_size * n_samples_per_prompt == global_batch_size.""" + product = _script_flag("rollout-batch-size") * _script_flag("n-samples-per-prompt") + assert product == _script_flag("global-batch-size"), ( + f"the example must whiten exactly one training batch; got {product} samples per rollout " + f"against a global batch of {_script_flag('global-batch-size')}" + ) + + +@pytest.mark.parametrize("doc", ["docs/zh/examples/algorithms.md", "docs/en/examples/algorithms.md"]) +def test_docs_do_not_describe_the_example_as_violating_the_boundary(doc): + """The known-deviation section used to cite the example as the + counterexample. + + It stopped being one when global_batch_size moved to 32, and the text was + left behind. Anything asserting num_rollout_minis is 2 for this example is + now false. + """ + import pathlib + + text = (pathlib.Path(__file__).resolve().parents[2] / doc).read_text(encoding="utf-8") + assert "num_rollout_minis = 2" not in text, f"{doc} still calls the shipped example a counterexample" diff --git a/tests/algorithms/test_gdpo.py b/tests/algorithms/test_gdpo.py new file mode 100644 index 000000000..8cc444028 --- /dev/null +++ b/tests/algorithms/test_gdpo.py @@ -0,0 +1,677 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""GDPO: per-reward group standardisation, weighted sum, batch whitening. + +The reference in ``_manual_gdpo`` implements arXiv 2601.05242 Eq. 4 and Eq. 7 +directly in plain Python, independently of the tensor implementation, so a +mistake in one is unlikely to be mirrored in the other. +""" + +import math +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + +from relax.algorithms import get_algorithm # noqa: E402 +from relax.algorithms.advantages import whiten_scalar # noqa: E402 +from relax.algorithms.rewards import REWARD_NORMALIZERS, extract_reward_components # noqa: E402 + + +def _args(keys=("correctness", "format"), weights=None, n=4): + return SimpleNamespace( + advantage_estimator="gdpo", + n_samples_per_prompt=n, + rewards_normalization=True, + grpo_std_normalization=True, + gdpo_reward_keys=list(keys) if keys is not None else None, + gdpo_reward_weights=weights, + ) + + +class _S: + """Minimal stand-in for Sample with the same reward-component contract.""" + + def __init__(self, group_index, reward): + self.group_index = group_index + self.reward = reward + + def get_reward_components(self, keys): + if not isinstance(self.reward, dict): + raise ValueError("Sample.reward must be a dict") + values = [] + for key in keys: + if key not in self.reward: + raise ValueError(f"Reward key {key!r} missing from sample reward") + values.append(self.reward[key]) + return values + + +def _normalize(args, samples): + spec = get_algorithm(args.advantage_estimator) + return REWARD_NORMALIZERS[spec.reward_normalizer](args, samples, [0.0] * len(samples)) + + +def _mk(groups, correctness, fmt): + return [_S(g, {"correctness": c, "format": f}) for g, c, f in zip(groups, correctness, fmt, strict=True)] + + +def _manual_gdpo(correctness, fmt, groups, weights=(1.0, 1.0)): + """Plain-Python reference for Eq. + + 4 + Eq. 7. + """ + per_key = [] + for column in (correctness, fmt): + out = [0.0] * len(column) + for g in sorted(set(groups)): + idx = [i for i, gg in enumerate(groups) if gg == g] + vals = [column[i] for i in idx] + mean = sum(vals) / len(vals) + var = sum((v - mean) ** 2 for v in vals) / (len(vals) - 1) + std = math.sqrt(var) + scale = max(abs(v) for v in vals) + collapsed = std <= 1e-6 * scale + # 1e-4, hardcoded on purpose: this oracle exists to pin parity with + # the reference implementation (trl grpo_trainer.py's scale_rewards + # GDPO path divides by std + 1e-4 at both steps), so importing our + # own constant here would make the test agree with itself. + for i in idx: + out[i] = 0.0 if collapsed else (column[i] - mean) / (std + 1e-4) + per_key.append(out) + return [weights[0] * a + weights[1] * b for a, b in zip(per_key[0], per_key[1], strict=True)] + + +# ---------------- registration ---------------- + + +def test_gdpo_is_registered(): + spec = get_algorithm("gdpo") + assert spec.reward_normalizer == "gdpo_decoupled" + assert spec.advantage_fn == "gdpo" + assert spec.policy_loss_fn == "ppo_clip" + assert spec.kl_level == "token" + + +def test_gdpo_spec_guards(): + spec = get_algorithm("gdpo") + assert spec.min_group_size == 2 + assert spec.allows_reward_post_process_hooks is False + assert spec.forbids_normalize_advantages is True + assert spec.requires_rewards_normalization is True + assert spec.uses_reward_components is True + + +# ---------------- step 1 + 2 ---------------- + + +def test_matches_hand_computed_reference(): + groups = [0, 0, 0, 0, 1, 1, 1, 1] + correctness = [1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0] + fmt = [1.0, 1.0, 0.0, 0.0, 0.5, 0.25, 0.75, 1.0] + samples = _mk(groups, correctness, fmt) + + got = _normalize(_args(), samples) + want = _manual_gdpo(correctness, fmt, groups) + + assert torch.allclose(torch.tensor(got), torch.tensor(want), atol=1e-6) + + +def test_weights_apply_to_normalized_advantages_not_raw_rewards(): + groups = [0, 0, 0, 0] + correctness = [1.0, 0.0, 1.0, 0.0] + fmt = [0.0, 0.0, 10.0, 10.0] + samples = _mk(groups, correctness, fmt) + + got = _normalize(_args(weights=[2.0, 0.5]), samples) + want = _manual_gdpo(correctness, fmt, groups, weights=(2.0, 0.5)) + + assert torch.allclose(torch.tensor(got), torch.tensor(want), atol=1e-6) + + +def test_component_scale_does_not_leak_into_the_combination(): + """After step 1 every component is unit-variance, so rescaling one is a no- + op.""" + groups = [0, 0, 0, 0] + correctness = [1.0, 0.0, 1.0, 0.0] + unit = [1.0, 2.0, 3.0, 4.0] + thousandfold = [1000.0, 2000.0, 3000.0, 4000.0] + + with_unit = _normalize(_args(), _mk(groups, correctness, unit)) + with_large = _normalize(_args(), _mk(groups, correctness, thousandfold)) + + # Not exactly a no-op: the additive epsilon does not scale with the data, so + # a 1000x rescale leaks about eps/std = 1e-4/1.291 = 7.7e-5. Measured 9.0e-5. + assert torch.allclose(torch.tensor(with_unit), torch.tensor(with_large), atol=1e-4) + + +def test_eps_makes_scale_invariance_approximate_for_tiny_rewards(): + """Documented limitation of the additive epsilon, shared with the GRPO + path. + + Dividing by ``std + eps`` is only scale-free while ``std >> eps``, and GDPO + uses ``eps = 1e-4`` to match the reference implementation rather than the + ``1e-6`` the GRPO path uses. That choice is not free: a component whose + spread is ~1e-3 is shrunk by roughly eps/std, which at 1e-4 is **7.2%** + against 0.08% at 1e-6. + + So a continuous reward with a very narrow spread -- the paper's maths setup + scores response length -- is damped noticeably more here than a reader + coming from the GRPO path would expect. Pinning the number is the point; + if it ever needs to be configurable, this is the test that says why. + """ + groups = [0, 0, 0, 0] + correctness = [1.0, 0.0, 1.0, 0.0] + + normal = _normalize(_args(), _mk(groups, correctness, [1.0, 2.0, 3.0, 4.0])) + tiny = _normalize(_args(), _mk(groups, correctness, [0.001, 0.002, 0.003, 0.004])) + + deviation = (torch.tensor(normal) - torch.tensor(tiny)).abs().max() + assert 0.07 < float(deviation) < 0.09, f"expected ~7% damping from eps=1e-4, got {float(deviation)}" + + +def test_default_weights_are_all_ones(): + groups = [0, 0, 0, 0] + samples = _mk(groups, [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + assert _normalize(_args(weights=None), samples) == _normalize(_args(weights=[1.0, 1.0]), samples) + + +def test_grouping_uses_group_index(): + correctness = [1.0, 0.0, 1.0, 0.0, 5.0, 5.0, 5.0, 5.0] + fmt = [1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0] + contiguous = _normalize(_args(), _mk([0, 0, 0, 0, 1, 1, 1, 1], correctness, fmt)) + # group 1 has collapsed correctness, so only format contributes there + assert abs(sum(contiguous[4:])) < 1e-5 + + +# ---------------- reward collapse ---------------- + + +def test_collapsed_component_contributes_zero_but_others_keep_signal(): + """This is GDPO's core benefit over GRPO.""" + groups = [0, 0, 0, 0] + correctness = [1.0, 1.0, 1.0, 1.0] # collapsed + fmt = [1.0, 0.0, 1.0, 0.0] + samples = _mk(groups, correctness, fmt) + + got = _normalize(_args(), samples) + fmt_only = _manual_gdpo([0.0] * 4, fmt, groups) + + assert torch.allclose(torch.tensor(got), torch.tensor(fmt_only), atol=1e-6) + assert max(abs(v) for v in got) > 0.5 + + +def test_fully_collapsed_group_yields_exact_zeros_without_nan(): + samples = _mk([0, 0, 0, 0], [0.7] * 4, [0.7] * 4) + got = _normalize(_args(), samples) + assert got == [0.0, 0.0, 0.0, 0.0] + + +def test_collapsed_group_does_not_leak_fp32_residual(): + """0.7 repeated 7x produces a nonzero residual if you rely on `/(std + eps)` alone.""" + samples = _mk([0] * 7, [0.7] * 7, [0.7] * 7) + got = _normalize(_args(n=7), samples) + assert got == [0.0] * 7 + + +@pytest.mark.parametrize("value", [0.0, 0.1, 0.7, 1.0, 1000.0, -3.5]) +def test_collapse_is_exact_at_any_magnitude(value): + samples = _mk([0] * 5, [value] * 5, [value] * 5) + assert _normalize(_args(n=5), samples) == [0.0] * 5 + + +def _grpo_reference(summed, groups): + """What GRPO does: sum the components first, then standardise once.""" + out = [0.0] * len(summed) + for g in sorted(set(groups)): + idx = [i for i, gg in enumerate(groups) if gg == g] + vals = [summed[i] for i in idx] + mean = sum(vals) / len(vals) + var = sum((v - mean) ** 2 for v in vals) / (len(vals) - 1) + std = math.sqrt(var) + for i in idx: + out[i] = 0.0 if std == 0 else (summed[i] - mean) / (std + 1e-4) + return out + + +def test_gdpo_distinguishes_reward_patterns_that_grpo_flattens(): + """The paper's core argument (arXiv 2601.05242, Sec. 3.1). + + With G=2, standardising collapses any two distinct values to +-1/sqrt(2), + so GRPO maps "one component fires" and "both components fire" to the exact + same advantage. Normalising per component first keeps the two apart. + """ + groups = [0, 0] + one_fires = _mk(groups, [0.0, 1.0], [0.0, 0.0]) + both_fire = _mk(groups, [0.0, 1.0], [0.0, 1.0]) + + grpo_one = _grpo_reference([0.0, 1.0], groups) + grpo_both = _grpo_reference([0.0, 2.0], groups) + assert torch.allclose(torch.tensor(grpo_one), torch.tensor(grpo_both), atol=1e-4), ( + "GRPO should be unable to tell these apart" + ) + + gdpo_one = _normalize(_args(n=2), one_fires) + gdpo_both = _normalize(_args(n=2), both_fire) + assert not torch.allclose(torch.tensor(gdpo_one), torch.tensor(gdpo_both), atol=1e-2) + # both_fire carries twice the signal: two components at +-1/sqrt(2) each + assert math.isclose(gdpo_both[1], 2 * gdpo_one[1], rel_tol=1e-4) + + +# ---------------- error contracts ---------------- + + +def test_missing_reward_key_raises(): + samples = [_S(0, {"correctness": 1.0}) for _ in range(4)] + with pytest.raises(ValueError, match="format"): + _normalize(_args(), samples) + + +def test_non_dict_reward_raises(): + samples = [_S(0, 1.0) for _ in range(4)] + with pytest.raises(ValueError, match="must be a dict"): + _normalize(_args(), samples) + + +def test_non_numeric_reward_raises(): + samples = [_S(0, {"correctness": 1.0, "format": "good"}) for _ in range(4)] + with pytest.raises(TypeError, match="must be a real number"): + _normalize(_args(), samples) + + +def test_bool_reward_raises(): + """bool subclasses int, so it would slip through a naive isinstance + check.""" + samples = [_S(0, {"correctness": True, "format": 1.0}) for _ in range(4)] + with pytest.raises(TypeError, match="must be a real number"): + _normalize(_args(), samples) + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) +def test_nan_or_inf_reward_raises(bad): + samples = [_S(0, {"correctness": bad, "format": 1.0}) for _ in range(4)] + with pytest.raises(ValueError, match="not finite"): + _normalize(_args(), samples) + + +def test_fewer_than_two_keys_raises(): + samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="at least two"): + _normalize(_args(keys=("correctness",)), samples) + + +def test_no_keys_raises(): + samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="at least two"): + _normalize(_args(keys=None), samples) + + +def test_duplicate_keys_raise(): + samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="duplicates"): + _normalize(_args(keys=("correctness", "correctness")), samples) + + +def test_weight_count_mismatch_raises(): + samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="gdpo-reward-weights"): + _normalize(_args(weights=[1.0]), samples) + + +def test_wrong_group_size_raises(): + samples = _mk([0, 0, 1, 1], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="expected 4"): + _normalize(_args(n=4), samples) + + +# ---------------- numerical properties ---------------- + + +def test_group_of_two_always_normalizes_to_plus_minus_one_over_sqrt_two(): + """With G=2 an unbiased std absorbs the magnitude entirely.""" + samples = _mk([0, 0], [0.0, 100.0], [0.0, 1.0]) + got = _normalize(_args(n=2), samples) + expected = 2 * (1.0 / math.sqrt(2.0)) + assert math.isclose(got[1], expected, rel_tol=1e-4) + assert math.isclose(got[0], -expected, rel_tol=1e-4) + + +def test_extract_reward_components_shape_and_dtype(): + samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + out = extract_reward_components(samples, ["correctness", "format"]) + assert out.shape == (4, 2) + assert out.dtype == torch.float32 + + +def test_extract_reward_components_preserves_key_order(): + samples = _mk([0, 0, 0, 0], [1.0, 2.0, 3.0, 4.0], [10.0, 20.0, 30.0, 40.0]) + forward = extract_reward_components(samples, ["correctness", "format"]) + reversed_ = extract_reward_components(samples, ["format", "correctness"]) + assert torch.equal(forward, reversed_.flip(dims=[1])) + + +def test_extract_accepts_python_ints(): + samples = [_S(0, {"correctness": 1, "format": 0}) for _ in range(2)] + out = extract_reward_components(samples, ["correctness", "format"]) + assert out.dtype == torch.float32 + + +def test_single_reward_gdpo_is_a_constant_positive_multiple_of_grpo(): + """Documented deviation: GDPO does NOT reduce to GRPO for one reward.""" + torch.manual_seed(0) + b, g = 16, 4 + rewards = (torch.rand(b * g) < 0.5).float().view(b, g) + grpo = ((rewards - rewards.mean(1, keepdim=True)) / (rewards.std(1, keepdim=True) + 1e-6)).flatten() + gdpo = whiten_scalar(grpo) + + mask = grpo.abs() > 1e-6 + ratio = gdpo[mask] / grpo[mask] + assert ratio.min() > 0 + assert torch.allclose(ratio, ratio[0].expand_as(ratio), atol=1e-4) + assert not torch.allclose(gdpo, grpo, atol=1e-3) + + +def test_sample_get_reward_components_matches_the_test_double(): + """The real Sample must honour the same contract as _S above.""" + from relax.utils.types import Sample + + sample = Sample(group_index=0, reward={"correctness": 1.0, "format": 0.5}) + assert sample.get_reward_components(["format", "correctness"]) == [0.5, 1.0] + + with pytest.raises(ValueError, match="missing from sample reward"): + sample.get_reward_components(["nope"]) + + scalar = Sample(group_index=0, reward=1.0) + with pytest.raises(ValueError, match="must be a dict"): + scalar.get_reward_components(["correctness"]) + + +def test_warns_once_when_every_component_collapses_in_every_group(caplog): + """A batch that produces no gradient at all must not be silent.""" + import logging + + samples = _mk([0, 0, 0, 0], [1.0] * 4, [0.5] * 4) + with caplog.at_level(logging.WARNING): + out = _normalize(_args(), samples) + + assert out == [0.0] * 4 + assert sum("all reward components collapsed" in r.message for r in caplog.records) == 1 + + +def test_does_not_warn_when_some_signal_survives(caplog): + import logging + + samples = _mk([0, 0, 0, 0], [1.0] * 4, [1.0, 0.0, 1.0, 0.0]) + with caplog.at_level(logging.WARNING): + _normalize(_args(), samples) + + assert not any("all reward components collapsed" in r.message for r in caplog.records) + + +def test_does_not_warn_when_only_some_groups_collapse(caplog): + import logging + + samples = _mk( + [0, 0, 0, 0, 1, 1, 1, 1], + [1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0], + [0.5, 0.5, 0.5, 0.5, 1.0, 0.0, 1.0, 0.0], + ) + with caplog.at_level(logging.WARNING): + _normalize(_args(), samples) + + assert not any("all reward components collapsed" in r.message for r in caplog.records) + + +# ---------------- reward value contract ---------------- + + +def test_numpy_scalars_are_accepted(): + """Reward functions routinely return numpy scalars; only float64 subclasses + float.""" + np = pytest.importorskip("numpy") + + for dtype in (np.float64, np.float32, np.int64, np.int32, np.float16): + samples = [ + _S(0, {"correctness": dtype(1), "format": dtype(0)}), + _S(0, {"correctness": dtype(0), "format": dtype(1)}), + ] + out = extract_reward_components(samples, ["correctness", "format"]) + assert out.dtype == torch.float32, dtype + assert out.shape == (2, 2), dtype + + +def test_numpy_bool_is_still_rejected(): + np = pytest.importorskip("numpy") + + samples = [_S(0, {"correctness": np.bool_(True), "format": 1.0}) for _ in range(4)] + with pytest.raises(TypeError, match="must be a real number"): + _normalize(_args(), samples) + + +def test_numpy_nan_is_still_rejected(): + np = pytest.importorskip("numpy") + + samples = [_S(0, {"correctness": np.float32("nan"), "format": 1.0}) for _ in range(4)] + with pytest.raises(ValueError, match="not finite"): + _normalize(_args(), samples) + + +# ---------------- weight contract ---------------- + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_weights_raise_instead_of_zeroing_the_batch(bad): + """Left unchecked these produce an all-zero batch with no error at all.""" + samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="not finite"): + _normalize(_args(weights=[bad, 1.0]), samples) + + +# ---------------- collapse detection is exact ---------------- + + +def test_component_with_small_relative_spread_is_kept(): + """A relative tolerance would have erased this component entirely.""" + groups = [0, 0, 0, 0] + correctness = [10000.0, 10000.005, 10000.010, 10000.015] + fmt = [1.0, 0.0, 1.0, 0.0] + got = _normalize(_args(), _mk(groups, correctness, fmt)) + + fmt_only = _manual_gdpo([0.0] * 4, fmt, groups) + assert not torch.allclose(torch.tensor(got), torch.tensor(fmt_only), atol=1e-3) + + +# ---------------- batch statistics must survive a large offset ---------------- +# +# distributed_mean_std used the one-pass E[x^2] - E[x]^2 form in the caller's +# float32. That subtracts two nearly equal large numbers whenever the values sit +# far from zero, and it failed silently in two different directions. Neither is +# hypothetical: the GDPO paper's maths setup uses a length reward, and token +# counts live exactly in this range. + + +def test_batch_std_survives_a_large_offset(): + """One-pass returned exactly 0 here, which reads as a collapsed batch.""" + from relax.algorithms.numerics import distributed_mean_std + + values = torch.tensor([1000.0, 1000.01, 1000.02, 1000.03]) + _mean, std = distributed_mean_std(values) + + assert std > 0, "a batch with real spread was reported as collapsed" + # rtol at float32 resolution, not float64: the statistic is computed in + # float64 but handed back in the caller's dtype. + torch.testing.assert_close(std.double(), values.double().std(), rtol=1e-6, atol=0) + + +def test_batch_std_is_not_inflated_by_a_large_offset(): + """One-pass returned 4.6 here against a true std of 1.3e-3.""" + from relax.algorithms.numerics import distributed_mean_std + + values = torch.tensor([10000.0, 10000.001, 10000.002, 10000.003]) + _mean, std = distributed_mean_std(values) + + torch.testing.assert_close(std.double(), values.double().std(), rtol=1e-6, atol=0) + assert std < 1e-2, f"std inflated to {float(std)}; advantages would be rescaled by ~1/{float(std):.3g}" + + +@pytest.mark.parametrize("offset", [0.0, 1e2, 1e3, 1e4]) +def test_whitening_is_shift_invariant(offset): + """Whitening centres before scaling, so adding a constant to every value + must not change the result. + + One-pass broke this well before float32 ran out of significand. + """ + from relax.algorithms.advantages import whiten_scalar + + base = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) + torch.testing.assert_close(whiten_scalar(base + offset), whiten_scalar(base), rtol=1e-5, atol=1e-6) + + +def test_batch_std_matches_torch_std_on_ordinary_input(): + """The fix must not move the numbers the existing algorithms already + produce.""" + from relax.algorithms.numerics import distributed_mean_std + + torch.manual_seed(20260726) + for n in (2, 4, 8, 64): + values = torch.randn(n) + mean, std = distributed_mean_std(values) + torch.testing.assert_close(mean, values.mean(), rtol=1e-6, atol=0) + torch.testing.assert_close(std, values.std(), rtol=1e-6, atol=0) + + +# ---------------- what step 3's boundary actually is ---------------- + + +def test_whitening_scope_is_whatever_the_caller_passes_not_a_training_batch(): + """Step 3 has no notion of a batch boundary; it whitens its argument. + + This is worth pinning because the surrounding docs used to claim the + statistics describe one training batch "matching the paper", and they do not. + The Megatron caller merges `num_rollout_minis` windows before calling + (actor.py: concat_rollout_batches, under the comment "we may need normalize + the whole rollout"), so with the shipped example's 4 * 8 against a + global_batch_size of 16 one whitening spans two optimizer steps. + + Concretely: whitening two batches together is not the same as whitening each. + """ + from relax.algorithms.advantages import whiten_scalar + + first = torch.tensor([1.0, 2.0, 3.0, 4.0]) + second = torch.tensor([101.0, 102.0, 103.0, 104.0]) + + merged = whiten_scalar(torch.cat([first, second])) + separate = torch.cat([whiten_scalar(first), whiten_scalar(second)]) + + assert not torch.allclose(merged, separate, atol=1e-3), ( + "if these agreed, the scope of step 3 would not matter and this test would be pointless" + ) + # The merged form is dominated by the between-batch offset, which is exactly + # the deviation from Eq. 6 the docs now describe. + assert merged[:4].max() < 0, "merged whitening puts the whole first batch below the mean" + assert torch.allclose(separate[:4], separate[4:], atol=1e-5), "per-batch whitening treats them alike" + + +# ---------------- step 3 now normalises per training batch ---------------- + + +def _gdpo_adv(rewards, mini_batch_sizes=None): + from relax.algorithms.advantages import compute_advantages_and_returns + + kl = [torch.zeros(1) for _ in rewards] + adv, _ = compute_advantages_and_returns( + SimpleNamespace(advantage_estimator="gdpo", kl_coef=0.0), + rewards=list(rewards), + kl=kl, + mini_batch_sizes=mini_batch_sizes, + ) + return torch.cat(adv) + + +def test_step_three_normalises_each_training_batch_separately(): + """Eq. + + 6's boundary. The caller merges num_rollout_minis batches before the + advantage stage, so without the counts one whitening covered all of them. + """ + from relax.algorithms.advantages import whiten_scalar + + first, second = [0.9, 1.1, 0.8, 1.2], [-1.2, -0.8, -1.1, -0.9] + got = _gdpo_adv(first + second, mini_batch_sizes=[4, 4]) + want = torch.cat([whiten_scalar(torch.tensor(first)), whiten_scalar(torch.tensor(second))]) + torch.testing.assert_close(got, want, rtol=1e-6, atol=1e-6) + + +def test_merging_the_batches_would_flip_signs_not_just_rescale(): + """Why the boundary matters: it decides which samples are reinforced. + + Whitening the two batches together centres both on the pooled mean, so + samples that were below their own batch's mean come out positive. Four of + these eight change sign -- this is a different objective, not a precision + difference. + """ + first, second = [0.9, 1.1, 0.8, 1.2], [-1.2, -0.8, -1.1, -0.9] + per_batch = _gdpo_adv(first + second, mini_batch_sizes=[4, 4]) + merged = _gdpo_adv(first + second, mini_batch_sizes=None) + + assert int((per_batch.sign() != merged.sign()).sum()) == 4 + + +def test_a_single_batch_is_unchanged_by_the_counts(): + """num_rollout_minis == 1 is the common case and must not move.""" + rewards = [1.0, 2.0, 3.0, 4.0] + torch.testing.assert_close(_gdpo_adv(rewards, [4]), _gdpo_adv(rewards, None)) + + +def test_counts_that_do_not_cover_the_shard_are_rejected(): + """Silently whitening the wrong window is the failure this replaces.""" + with pytest.raises(ValueError, match="sum to"): + _gdpo_adv([1.0, 2.0, 3.0, 4.0], mini_batch_sizes=[3, 3]) + + +@pytest.mark.parametrize("estimator", ["grpo", "gspo", "sapo", "cispo"]) +def test_other_estimators_ignore_the_batch_counts(estimator): + """They absorb it in **_unused, so passing it must be bit-identical.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + args = SimpleNamespace(advantage_estimator=estimator, kl_coef=0.0) + inputs = dict(rewards=[1.0, -1.0, 0.5, -0.5], kl=[torch.zeros(2) for _ in range(4)]) + + without, _ = compute_advantages_and_returns(args, **inputs) + with_counts, _ = compute_advantages_and_returns(args, mini_batch_sizes=[2, 2], **inputs) + + for left, right in zip(without, with_counts, strict=True): + assert torch.equal(left, right) + + +def test_all_zero_weights_are_rejected(): + """Every component times zero is a batch of zero advantages and a clean + exit.""" + from relax.algorithms.rewards import resolve_gdpo_weights + + with pytest.raises(ValueError, match="all zero"): + resolve_gdpo_weights(_args(weights=[0.0, 0.0]), ["correctness", "format"]) + + +def test_one_zero_weight_is_allowed(): + """Muting a component is a legitimate configuration; only muting all is + not.""" + from relax.algorithms.rewards import resolve_gdpo_weights + + assert resolve_gdpo_weights(_args(weights=[1.0, 0.0]), ["correctness", "format"]) == [1.0, 0.0] + + +@pytest.mark.parametrize("bad", [[], [0, 4], [2.5, 1.5], [-1, 5]]) +def test_malformed_batch_sizes_raise_rather_than_falling_back(bad): + """An empty or malformed list must not quietly restore merged whitening.""" + with pytest.raises(ValueError, match="positive ints"): + _gdpo_adv([1.0, 2.0, 3.0, 4.0], mini_batch_sizes=bad) + + +def test_a_single_segment_is_still_size_checked(): + """[4] on a 5-sample shard is a real mismatch, not a request to use one + window.""" + with pytest.raises(ValueError, match="sum to"): + _gdpo_adv([1.0, 2.0, 3.0, 4.0, 5.0], mini_batch_sizes=[4]) diff --git a/tests/algorithms/test_policy_loss_dispatch.py b/tests/algorithms/test_policy_loss_dispatch.py index e0934178e..c8fe6ecf2 100644 --- a/tests/algorithms/test_policy_loss_dispatch.py +++ b/tests/algorithms/test_policy_loss_dispatch.py @@ -108,7 +108,7 @@ def test_sapo_taus_are_read_from_args(): assert torch.equal(got[0], want[0]) -@pytest.mark.parametrize("estimator", ["grpo", "gspo", "ppo", "reinforce_plus_plus"]) +@pytest.mark.parametrize("estimator", ["grpo", "gspo", "gdpo", "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) diff --git a/tests/algorithms/test_post_process_rewards_dispatch.py b/tests/algorithms/test_post_process_rewards_dispatch.py index 64c5b6003..cd1fa2d6a 100644 --- a/tests/algorithms/test_post_process_rewards_dispatch.py +++ b/tests/algorithms/test_post_process_rewards_dispatch.py @@ -139,8 +139,8 @@ 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. + GDPO run), 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}), From 905d0e06b6ab690163cc103192e9e7b4c03508b0 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 13:25:38 +0800 Subject: [PATCH 03/22] fix(gdpo): agree on the segmentation before whitening, reject float32 overflow Two failure modes that were silent rather than loud. Segmentation. `_whiten_by_segment` runs one collective per segment, and the comment claimed every rank would agree on the count because `num_rollout_minis` comes from the minibatch plan. Expected is not checked, and the consequence of being wrong is a deadlock: a rank expecting two segments and a rank expecting three both block on the third collective, with no traceback and no exit. The same applies to a rank whose `mini_batch_sizes` is malformed -- raising locally strands its peers inside the collective sequence, which is worse than not validating at all. Both verdicts now travel in one MAX all-reduce ahead of any segment, so every rank reads the same numbers and raises together. Tests cover unequal shards across ranks (a DP split balances tokens, not sample counts), mismatched counts, the `None`-versus-segmented case, and malformed metadata on one rank; each asserts that *both* ranks return, so a regression times out instead of passing quietly. Overflow. `extract_reward_components` checked `math.isfinite` on the float64 value, then cast to float32. A reward of 1e300 passes that check and becomes `inf`, which reads as a non-finite std one stage later, zeroes the batch, and lets the run finish cleanly having trained on nothing. Rewards above float32's range are now rejected at the boundary, where the message can name the reward function that produced them. 1e30 still passes: the bound is the representable range, not an opinion about reward magnitude. Tests: 1645 passed, same 2 failures + 2 errors as main@98a1274 here. --- relax/algorithms/advantages.py | 83 +++++++++++++--- relax/algorithms/rewards.py | 25 ++++- .../algorithms/test_distributed_whitening.py | 98 +++++++++++++++++++ tests/algorithms/test_gdpo.py | 20 ++++ 4 files changed, 214 insertions(+), 12 deletions(-) diff --git a/relax/algorithms/advantages.py b/relax/algorithms/advantages.py index 781a6dc08..99fd755be 100644 --- a/relax/algorithms/advantages.py +++ b/relax/algorithms/advantages.py @@ -79,26 +79,87 @@ def advantage_grpo_broadcast(args: Any, *, rewards, kl, **_unused): return advantages, returns +def _agree_on_segmentation(n_segments: int, local_error: str, values, process_group) -> None: + """Reach one verdict on the segmentation across the whole group. + + Each segment costs one collective, so this has to agree *before* any of + them run. Two ways it can disagree, both of which deadlock rather than + fail: + + * The segment counts differ. A rank expecting two segments and a rank + expecting three both block on the third collective -- no traceback, no + exit code, just a job that never finishes. ``num_rollout_minis`` comes + from the minibatch plan and is expected to be uniform, but "expected" is + not "checked". + * One rank's ``mini_batch_sizes`` is malformed. Raising locally is worse + than not checking at all: that rank leaves the collective sequence while + every other rank is still waiting inside it. + + So the local verdict travels *with* the count, in one MAX reduction (the + low bound negated, the way :func:`~relax.algorithms.numerics.is_collapsed` + does it). Every rank reads the same three numbers and therefore raises + together, at the same call, with a message naming which failure it was. + """ + if process_group is None: + if local_error: + raise ValueError(local_error) + return + + # A rank that already failed contributes 0 segments: a neutral value that + # cannot be mistaken for a real count, and one that trips the mismatch + # branch too if the `any_bad` branch were ever removed. + flags = torch.tensor( + [-n_segments if not local_error else 0, n_segments if not local_error else 0, 1 if local_error else 0], + dtype=torch.int64, + device=values.device, + ) + dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=process_group) + low, high, any_bad = -int(flags[0]), int(flags[1]), int(flags[2]) + + if any_bad: + raise ValueError( + local_error + or "another rank reported malformed mini_batch_sizes; every rank fails here so the " + "per-segment collectives cannot deadlock." + ) + if low != high: + raise ValueError( + f"mini_batch_sizes describes {n_segments} segment(s) on this rank, but ranks in the group " + f"report between {low} and {high}. Every rank must whiten the same number of segments, " + "or the per-segment collectives deadlock." + ) + + def _whiten_by_segment(values, mini_batch_sizes, process_group): """Whiten each training batch separately, in the order they were merged. - Every rank runs the same number of segments -- ``num_rollout_minis`` comes - from the minibatch plan rather than from the data -- so the per-segment - collectives stay matched across the data-parallel group. + ``num_rollout_minis`` comes from the minibatch plan rather than from the + data, so every rank is expected to run the same number of segments; + :func:`_agree_on_segmentation` is what turns that expectation into a + checked precondition rather than a deadlock. """ - if mini_batch_sizes is None: - return whiten_scalar(values, process_group=process_group) # Validate whatever was passed, including a single segment: treating an empty # or malformed list as "fall back to one window" would silently restore the - # merged behaviour this function exists to replace. - if not mini_batch_sizes or any(not isinstance(n, int) or n <= 0 for n in mini_batch_sizes): - raise ValueError(f"mini_batch_sizes must be a non-empty list of positive ints, got {mini_batch_sizes}.") - if sum(mini_batch_sizes) != values.numel(): - raise ValueError( + # merged behaviour this function exists to replace. The verdict is not acted + # on until the whole group has shared it. + local_error = "" + if mini_batch_sizes is None: + n_segments = 1 + elif not mini_batch_sizes or any(not isinstance(n, int) or n <= 0 for n in mini_batch_sizes): + n_segments = 0 + local_error = f"mini_batch_sizes must be a non-empty list of positive ints, got {mini_batch_sizes}." + elif sum(mini_batch_sizes) != values.numel(): + n_segments = 0 + local_error = ( f"mini_batch_sizes {mini_batch_sizes} sum to {sum(mini_batch_sizes)}, " f"but this rank holds {values.numel()} samples." ) - if len(mini_batch_sizes) == 1: + else: + n_segments = len(mini_batch_sizes) + + _agree_on_segmentation(n_segments, local_error, values, process_group) + + if mini_batch_sizes is None or n_segments == 1: return whiten_scalar(values, process_group=process_group) out, start = [], 0 for size in mini_batch_sizes: diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index e1c28a6fe..6e621470b 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -25,6 +25,10 @@ GROUP_EPS = STD_EPS +_FLOAT32_MAX = float(torch.finfo(torch.float32).max) +"""Largest finite float32. Rewards are carried as float32 from here on, so a +value above this is an overflow waiting to happen rather than a large reward.""" + def group_positions(samples: list[Any], expected_size: int) -> dict[int, list[int]]: """Map ``Sample.group_index`` to the positions it occupies in ``samples``. @@ -193,9 +197,28 @@ def extract_reward_components(samples: list[Any], keys: list[str]) -> torch.Tens numeric = float(value) if not math.isfinite(numeric): raise ValueError(f"Reward {key!r} of sample {position} is {numeric}, which is not finite.") + # Finite in float64 is not enough: the tensor below is float32, whose + # largest value is ~3.4e38. A reward of 1e300 passes `isfinite` here, + # becomes `inf` on cast, and then reads as a non-finite std one stage + # later -- where the batch is zeroed and the run trains on no signal + # without ever failing. Catching the overflow at the boundary keeps + # the diagnosis at the reward function that produced it. + if abs(numeric) > _FLOAT32_MAX: + raise ValueError( + f"Reward {key!r} of sample {position} is {numeric!r}, which overflows float32 " + f"(max {_FLOAT32_MAX:.6g}). Rescale the reward; casting it would silently produce inf." + ) row.append(numeric) rows.append(row) - return torch.tensor(rows, dtype=torch.float32) + + components = torch.tensor(rows, dtype=torch.float32) + # Belt and braces: the per-value check above is exact, but it only sees what + # `float()` returned. Anything that slips past it must not reach the + # normaliser, where non-finite input is indistinguishable from a collapse. + if not torch.isfinite(components).all(): + bad = (~torch.isfinite(components)).nonzero()[0].tolist() + raise ValueError(f"Reward {keys[bad[1]]!r} of sample {bad[0]} is not finite after casting to float32.") + return components def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: diff --git a/tests/algorithms/test_distributed_whitening.py b/tests/algorithms/test_distributed_whitening.py index 7a438e32f..fb71a4798 100644 --- a/tests/algorithms/test_distributed_whitening.py +++ b/tests/algorithms/test_distributed_whitening.py @@ -24,6 +24,39 @@ SHARDS = [[-0.7, 0.7], [-1.4, 1.4]] +def _t(values): + return torch.tensor(values, dtype=torch.float32) + + +# Per-rank (values, mini_batch_sizes) for the segmented cases. The shard sizes +# differ between ranks on purpose: a data-parallel split is balanced by tokens, +# not by sample count, so equal-sized shards are the easy case rather than the +# real one. +_SEGMENT_MODES = { + # Two segments on both ranks, split 2+2 on rank 0 and 1+3 on rank 1. + # Segment 0 spans [1, 3 | 5] and segment 1 spans [10, 30 | 50, 70, 90], so a + # per-segment statistic and a whole-shard statistic give different answers. + "segmented_uneven": lambda rank: ( + (_t([1.0, 3.0, 10.0, 30.0]), [2, 2]) if rank == 0 else (_t([5.0, 50.0, 70.0, 90.0]), [1, 3]) + ), + # Rank 0 wants two segments, rank 1 wants one. Two collectives against one: + # the deadlock this check exists to prevent. + "segment_count_mismatch": lambda rank: ( + (_t([1.0, 3.0, 10.0, 30.0]), [2, 2]) if rank == 0 else (_t([5.0, 50.0, 70.0, 90.0]), [4]) + ), + # Rank 1's metadata does not describe its shard. Raising locally would strand + # rank 0 inside the collective sequence. + "malformed_on_one_rank": lambda rank: ( + (_t([1.0, 3.0, 10.0, 30.0]), [2, 2]) if rank == 0 else (_t([5.0, 50.0, 70.0, 90.0]), [2, 99]) + ), + # Rank 0 passes no segmentation at all (one whole-shard window) while rank 1 + # asks for two. Same deadlock, reached through the `None` branch. + "none_versus_segmented": lambda rank: ( + (_t([1.0, 3.0, 10.0, 30.0]), None) if rank == 0 else (_t([5.0, 50.0, 70.0, 90.0]), [2, 2]) + ), +} + + def _run(rank, world_size, port, mode, out): os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = str(port) @@ -52,6 +85,18 @@ def _run(rank, world_size, port, mode, out): # Each shard is constant on its own but the batch is not. values = torch.tensor([0.7, 0.7] if rank == 0 else [1.4, 1.4], dtype=torch.float32) out[rank] = [is_collapsed(values, process_group=group)] + elif mode in _SEGMENT_MODES: + from relax.algorithms.advantages import _whiten_by_segment + + values, sizes = _SEGMENT_MODES[mode](rank) + try: + out[rank] = _whiten_by_segment(values, sizes, group).tolist() + except ValueError as exc: + # Recorded rather than raised: the point of these cases is that + # *both* ranks come back, so a hang is distinguishable from a + # rejection. A propagating exception would look the same as a + # rank that never reached the collective. + out[rank] = f"ValueError: {exc}" else: # pragma: no cover - guard against typos in the test itself raise AssertionError(mode) finally: @@ -100,3 +145,56 @@ def test_collapse_is_decided_globally_not_per_shard(): per_shard = _spawn("collapsed_only_locally") assert per_shard == {0: [False], 1: [False]}, "each shard is constant, the batch is not" + + +# ---------------- segmented whitening (GDPO step 3 across merged batches) ---------------- + + +def test_each_segment_is_whitened_against_the_whole_group(): + """Per-segment statistics, reduced across ranks, on unequal shards. + + Segment 0 holds [1, 3] on rank 0 and [5] on rank 1; segment 1 holds [10, + 30] and [50, 70, 90]. Whitening each segment against the group means each + segment's *joint* mean is 0 -- neither rank's own slice is centred, which + is what separates this from shard-local whitening. + """ + out = _spawn("segmented_uneven") + assert isinstance(out[0], list) and isinstance(out[1], list), out + + seg0 = torch.tensor(out[0][:2] + out[1][:1]) + seg1 = torch.tensor(out[0][2:] + out[1][1:]) + assert abs(seg0.mean().item()) < 1e-5, f"segment 0 not centred across ranks: {seg0}" + assert abs(seg1.mean().item()) < 1e-5, f"segment 1 not centred across ranks: {seg1}" + + # And the two segments really were separate windows: whitening the merged + # shard instead would leave segment 0 (values 1-5) far below segment 1 + # (values 10-90) rather than both centred on 0. + assert seg0.std().item() > 0.5, "segment 0 collapsed; it should carry its own spread" + assert seg1.std().item() > 0.5, "segment 1 collapsed; it should carry its own spread" + + +def test_mismatched_segment_counts_fail_on_every_rank(): + """The deadlock case: unequal collective counts must raise, not hang. + + Both ranks returning at all is the assertion. If the check were removed, + this test would time out rather than fail. + """ + out = _spawn("segment_count_mismatch") + for rank in (0, 1): + assert isinstance(out[rank], str), f"rank {rank} did not raise: {out[rank]!r}" + assert "same number of segments" in out[rank], out[rank] + + +def test_no_segmentation_on_one_rank_is_also_a_mismatch(): + out = _spawn("none_versus_segmented") + for rank in (0, 1): + assert isinstance(out[rank], str), f"rank {rank} did not raise: {out[rank]!r}" + assert "same number of segments" in out[rank], out[rank] + + +def test_malformed_metadata_on_one_rank_fails_both(): + """A local raise would strand the other rank inside the collectives.""" + out = _spawn("malformed_on_one_rank") + assert isinstance(out[1], str) and "sum to" in out[1], out[1] + assert isinstance(out[0], str), f"rank 0 did not fail with its peer: {out[0]!r}" + assert "another rank reported malformed" in out[0], out[0] diff --git a/tests/algorithms/test_gdpo.py b/tests/algorithms/test_gdpo.py index 8cc444028..e72884af8 100644 --- a/tests/algorithms/test_gdpo.py +++ b/tests/algorithms/test_gdpo.py @@ -296,6 +296,26 @@ def test_nan_or_inf_reward_raises(bad): _normalize(_args(), samples) +@pytest.mark.parametrize("huge", [1e300, -1e300, 3.5e38]) +def test_reward_that_overflows_float32_raises_rather_than_becoming_inf(huge): + """Finite in float64, `inf` after the cast — the gap `isfinite` misses. + + Left unchecked this is silent, not loud: the cast produces `inf`, the group + std reads as non-finite one stage later, `whiten_scalar` returns zeros, and + the run trains on no signal while exiting cleanly. + """ + samples = [_S(0, {"correctness": huge, "format": 1.0}) for _ in range(4)] + with pytest.raises(ValueError, match="overflows float32"): + _normalize(_args(), samples) + + +def test_a_large_but_representable_reward_is_accepted(): + """The bound is float32's range, not an opinion about reward magnitude.""" + samples = _mk([0, 0, 0, 0], [1e30, 2e30, 3e30, 4e30], [1.0, 0.0, 1.0, 0.0]) + out = _normalize(_args(), samples) + assert all(math.isfinite(v) for v in out), out + + def test_fewer_than_two_keys_raises(): samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) with pytest.raises(ValueError, match="at least two"): From e5404dc8407d08bf6205a19313462a6d9129fc80 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 14:16:26 +0800 Subject: [PATCH 04/22] feat(gdpo): make the upstream reward consumers component-aware Three consumers sit ahead of the advantage stage and ask "did this prompt group carry signal?" -- the dynamic-sampling filter, and the zero-std metrics on both the agentic and the Ray rollout paths. All three answered it with the single scalar `--reward-key` selects. For a multi-reward algorithm that is the wrong question, and wrong in the direction that undoes the algorithm. A group of `(1, 0)` and `(0, 1)` rollouts has an identical summed reward on every sample, so the filter reads it as dead and drops it before training -- while GDPO would have extracted signal from both components. Those are precisely the groups the estimator exists to keep, so leaving this unfixed means the reward stage does its job and the sampler throws the result away. `group_carries_reward_signal` asks the question the configured algorithm actually needs: the `--reward-key` scalar for single-reward algorithms, every component for multi-reward ones (a group is dead only when all of them are flat). The three call sites now share it. Two incidental fixes that fall out of touching these lines: - The zero-std metrics gated on `advantage_estimator == "ppo"`, two more hard-coded algorithm names. They now read `needs_critic`, which is what the gate meant: value-based estimators do not compute group-relative rewards. Same set of algorithms, no behaviour change. - The filter used `std > 0` while the metrics used exact equality, for the same question. Unified on exact equality, matching `is_collapsed` and for the reason argued there. The two differ only for a group that is exactly flat yet whose computed standard deviation is not exactly 0. The `--dynamic-sampling-filter-path` warning now fires only for a *custom* filter; the built-in one is correct here and no longer warrants one. Also fixes a flake I introduced in the previous commit: the segmented whitening tests derived a fixed port from `hash(mode)`, which is salted per interpreter and collided with `test_gdn_cp_reassembly`'s dynamically chosen port. Both now ask the OS for a free port. Tests: 1653 passed, twice in a row; same 2 failures + 2 errors as main@98a1274 here. --- relax/agentic/rollout.py | 22 +-- relax/algorithms/rewards.py | 36 +++++ relax/distributed/ray/rollout.py | 12 +- .../filters/dynamic_sampling_filters.py | 16 ++- relax/utils/arguments.py | 22 +-- .../algorithms/test_distributed_whitening.py | 19 ++- .../algorithms/test_multi_reward_consumers.py | 136 ++++++++++++++++++ 7 files changed, 234 insertions(+), 29 deletions(-) create mode 100644 tests/algorithms/test_multi_reward_consumers.py diff --git a/relax/agentic/rollout.py b/relax/agentic/rollout.py index 622fa5eb7..ec37e392a 100644 --- a/relax/agentic/rollout.py +++ b/relax/agentic/rollout.py @@ -19,6 +19,8 @@ get_agentic_runtime_resources, ) from relax.agentic.profile import TRACE_KEY +from relax.algorithms import get_algorithm +from relax.algorithms.rewards import group_carries_reward_signal from relax.engine.filters.base_types import MetricGatherer, call_dynamic_filter from relax.engine.rollout import on_policy_distillation as opd from relax.engine.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput @@ -1293,19 +1295,19 @@ def _dict_add_prefix(d, prefix): def _compute_zero_std_metrics(args, all_samples: list[Sample]) -> dict[str, float]: - if args.advantage_estimator == "ppo": + if get_algorithm(args.advantage_estimator).needs_critic: return {} all_sample_groups = group_by(all_samples, lambda sample: sample.group_index) - reward_groups = [ - [sample.get_reward_value(args) for sample in group if sample.reward is not None] - for group in all_sample_groups.values() - ] - interesting_rewards = [ - str(round(rewards[0], 1)) - for rewards in reward_groups - if rewards and all(rewards[0] == reward for reward in rewards) - ] + interesting_rewards = [] + for group in all_sample_groups.values(): + rewarded = [sample for sample in group if sample.reward is not None] + # Counted as flat only when it carries no signal *for this algorithm*: + # for a multi-reward one that means every component is flat, not that + # the --reward-key scalar happens to be. + if not rewarded or group_carries_reward_signal(args, rewarded): + continue + interesting_rewards.append(str(round(rewarded[0].get_reward_value(args), 1))) return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index 6e621470b..f04d00855 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -17,6 +17,7 @@ import torch from relax.algorithms.numerics import GDPO_EPS, STD_EPS, collapsed_columns +from relax.algorithms.spec import get_algorithm from relax.utils.logging_utils import get_logger from relax.utils.training.ppo_utils import compute_rloo_leave_one_out_rewards @@ -271,6 +272,41 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl return (normalized * weight_tensor).sum(dim=1).tolist() +def group_carries_reward_signal(args: Any, samples: list[Any]) -> bool: + """Whether one prompt group still carries signal for *this* algorithm. + + Consumers upstream of the advantage stage -- the dynamic-sampling filter, + the zero-std metrics -- ask "did this group vary?" to decide whether it is + worth keeping or worth reporting. They have always answered it with the + single scalar ``--reward-key`` selects, which is the right question only + for an algorithm that consumes that scalar. + + A multi-reward algorithm standardises each component separately, so a group + is dead only when *every* component is flat. Judging GDPO by the summed + scalar throws away exactly the groups it exists to keep: ``(1, 0)`` and + ``(0, 1)`` have identical sums and different components. + + The two consumers did not previously agree on the single-reward test: + ``check_reward_nonzero_std`` used ``std > 0``, the zero-std metrics used + exact equality. This unifies them on exact equality, matching + :func:`relax.algorithms.numerics.is_collapsed` and for the same reason -- + a tolerance wide enough to absorb float error also discards real signal, + and the two answers differ only for a group that is exactly flat yet whose + computed standard deviation is not exactly 0. + """ + if not samples: + return False + + spec = get_algorithm(args.advantage_estimator) + if not spec.uses_reward_components: + rewards = [sample.get_reward_value(args) for sample in samples] + return any(reward != rewards[0] for reward in rewards) + + keys = resolve_gdpo_keys(args) + components = extract_reward_components(samples, keys) + return bool((~collapsed_columns(components, dim=0)).any()) + + REWARD_NORMALIZERS: dict[str, Callable[[Any, list[Any], list[float]], list[float]]] = { "none": normalize_none, "group_mean": normalize_group_mean, diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index db8255b29..776f10211 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -21,6 +21,8 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS +from relax.algorithms import get_algorithm +from relax.algorithms.rewards import group_carries_reward_signal from relax.backends.sglang.sglang_engine import SGLangEngine from relax.distributed.ray.rollout_validation import validate_server_group_gpu_indices from relax.engine.rollout.base_types import call_rollout_fn @@ -4051,12 +4053,16 @@ def token_perf(response_lengths, non_generation_time, key=""): def _compute_zero_std_metrics(args, all_samples: list[Sample]): # only compute in GRPO-like algorithms where one prompt has multiple responses - if args.advantage_estimator == "ppo": + if get_algorithm(args.advantage_estimator).needs_critic: return {} def _is_zero_std(samples: list[Sample]): - rewards = [sample.get_reward_value(args) for sample in samples] - return len(rewards) == 0 or all(rewards[0] == r for r in rewards) + # Reads whichever notion of "signal" this algorithm uses: the + # --reward-key scalar for single-reward algorithms, every component for + # multi-reward ones. Counting a GDPO group as zero-std because its + # summed reward is flat overstates the count and reads as "most of the + # batch is dead" when it is not. + return not group_carries_reward_signal(args, samples) all_sample_groups = group_by(all_samples, lambda s: s.group_index) interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)] diff --git a/relax/engine/filters/dynamic_sampling_filters.py b/relax/engine/filters/dynamic_sampling_filters.py index 0b50cd8f3..501e352df 100644 --- a/relax/engine/filters/dynamic_sampling_filters.py +++ b/relax/engine/filters/dynamic_sampling_filters.py @@ -1,7 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -import torch - +from relax.algorithms.rewards import group_carries_reward_signal from relax.engine.filters.base_types import DynamicFilterOutput from relax.utils.types import Sample @@ -10,9 +9,16 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): - rewards = [sample.get_reward_value(args) for sample in samples] - keep = torch.tensor(rewards, dtype=torch.float).std() > 0.0 + """Drop prompt groups whose rewards carry no signal for this algorithm. + + "No signal" is algorithm-dependent, which is why the test is not spelled + out here. For a single-reward algorithm it is the standard deviation of the + ``--reward-key`` scalar, exactly as before. For a multi-reward algorithm it + is whether *every* component is flat -- judging those by the summed scalar + would drop the groups the algorithm exists to keep. + """ + keep = group_carries_reward_signal(args, samples) return DynamicFilterOutput( keep=keep, - reason=None if keep else f"zero_std_{round(rewards[0], 1)}", + reason=None if keep else f"zero_std_{round(samples[0].get_reward_value(args), 1)}", ) diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index b1d4fe54f..95117ee77 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -3100,17 +3100,21 @@ def _validate_multi_reward_args(args, spec) -> None: "used for metrics and for the raw_reward column." ) - if args.dynamic_sampling_filter_path: - # A warning rather than an error: the filter is opt-in and a custom one may - # well be component-aware. The built-in check_reward_nonzero_std is not — it - # reads the single --reward-key scalar, so a group where that component is - # flat but another still varies gets dropped, which is exactly the case this - # estimator exists to keep. + filter_path = args.dynamic_sampling_filter_path + if filter_path and not filter_path.endswith("check_reward_nonzero_std"): + # The built-in filter reads `group_carries_reward_signal`, so it judges a + # group by every component and needs no warning. A custom one is opaque + # from here: if it reduces the group to the single --reward-key scalar, + # it drops exactly the groups this estimator exists to keep. A warning + # rather than an error, because a custom filter may well be + # component-aware and we cannot tell. logger.warning( - "%r combines multiple reward components, but --dynamic-sampling-filter-path filters on the " - "single --reward-key scalar (%r). Groups carrying signal only in the other components may be " - "dropped before training sees them.", + "%r combines multiple reward components, but --dynamic-sampling-filter-path points at a " + "custom filter (%s). If it judges a group by the single --reward-key scalar (%r), groups " + "carrying signal only in the other components will be dropped before training sees them. " + "relax.engine.filters.dynamic_sampling_filters.check_reward_nonzero_std handles this correctly.", spec.name, + filter_path, args.reward_key, ) diff --git a/tests/algorithms/test_distributed_whitening.py b/tests/algorithms/test_distributed_whitening.py index fb71a4798..3137c3bef 100644 --- a/tests/algorithms/test_distributed_whitening.py +++ b/tests/algorithms/test_distributed_whitening.py @@ -103,11 +103,26 @@ def _run(rank, world_size, port, mode, out): dist.destroy_process_group() +def _free_port() -> int: + """Ask the OS for a port instead of deriving one from the mode name. + + A fixed port per mode collides with any other test that happens to pick the + same one, and `hash()` on a str is salted per interpreter, so which port a + mode gets changes between runs -- the collision is intermittent and lands on + whichever test ran second. `tests/backends/megatron/test_gdn_cp_reassembly.py` + already does it this way. + """ + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + def _spawn(mode, world_size=2): manager = mp.Manager() out = manager.dict() - port = 29500 + abs(hash(mode)) % 2000 - mp.spawn(_run, args=(world_size, port, mode, out), nprocs=world_size, join=True) + mp.spawn(_run, args=(world_size, _free_port(), mode, out), nprocs=world_size, join=True) return dict(out) diff --git a/tests/algorithms/test_multi_reward_consumers.py b/tests/algorithms/test_multi_reward_consumers.py new file mode 100644 index 000000000..3ab588531 --- /dev/null +++ b/tests/algorithms/test_multi_reward_consumers.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Consumers upstream of the advantage stage must ask the algorithm's question. + +The dynamic-sampling filter and the zero-std metrics both decide whether a +prompt group "carries signal". They answered that with the single +``--reward-key`` scalar, which is only the right question for an algorithm that +consumes that scalar. + +For a multi-reward algorithm it is the wrong one, and wrong in the direction +that undoes the algorithm: a group where ``correctness`` is flat but ``format`` +still varies looks dead by the summed scalar, gets dropped by the filter, and +never reaches training -- while GDPO would have extracted signal from the +component that did vary. Keeping those groups is the entire point. +""" + +from types import SimpleNamespace + +import pytest + + +torch = pytest.importorskip("torch") + +from relax.algorithms.rewards import group_carries_reward_signal # noqa: E402 + + +class _S: + """Minimal stand-in for Sample with the same reward contract.""" + + 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 get_reward_components(self, keys): + return [self.reward[key] for key in keys] + + +def _gdpo_args(n=4): + return SimpleNamespace( + advantage_estimator="gdpo", + n_samples_per_prompt=n, + reward_key="score", + gdpo_reward_keys=["correctness", "format"], + gdpo_reward_weights=None, + ) + + +def _grpo_args(): + return SimpleNamespace(advantage_estimator="grpo", n_samples_per_prompt=4, reward_key="score") + + +def _group(correctness, fmt): + """One prompt group; `score` is the summed scalar --reward-key selects.""" + return [_S(0, {"correctness": c, "format": f, "score": c + f}) for c, f in zip(correctness, fmt, strict=True)] + + +# ---------------- the case the whole feature exists for ---------------- + + +def test_group_flat_in_one_component_still_carries_signal(): + """`correctness` constant, `format` varying: one live component is enough. + + The scalar happens to vary here too, so this case alone does not separate + the two tests -- it pins the weaker property that a partially collapsed + group is not treated as dead. The case that does separate them is below. + """ + group = _group(correctness=[1.0, 1.0, 1.0, 1.0], fmt=[0.0, 0.5, 1.0, 0.5]) + assert group_carries_reward_signal(_gdpo_args(), group) is True + + +def test_equal_sums_from_different_components_carry_signal(): + """(1,0) and (0,1) both sum to 1 -- the paper's motivating case. + + Every sample's --reward-key scalar is identical, so this is precisely the + group a scalar-based filter throws away. + """ + group = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[0.0, 1.0, 0.0, 1.0]) + scalars = {s.get_reward_value(_gdpo_args()) for s in group} + assert scalars == {1.0}, "precondition: the scalar really is constant" + + assert group_carries_reward_signal(_gdpo_args(), group) is True + + +def test_group_flat_in_every_component_is_dead(): + """GDPO does not invent signal: all components constant means no signal.""" + group = _group(correctness=[1.0, 1.0, 1.0, 1.0], fmt=[0.5, 0.5, 0.5, 0.5]) + assert group_carries_reward_signal(_gdpo_args(), group) is False + + +# ---------------- single-reward algorithms keep their old answer ---------------- + + +def test_single_reward_algorithm_reads_only_the_reward_key(): + varying = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[0.0, 1.0, 0.0, 1.0]) + # Components vary, but the scalar GRPO consumes does not. + assert group_carries_reward_signal(_grpo_args(), varying) is False + + real = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[1.0, 0.0, 1.0, 0.0]) + assert group_carries_reward_signal(_grpo_args(), real) is True + + +def test_empty_group_carries_nothing(): + assert group_carries_reward_signal(_gdpo_args(), []) is False + assert group_carries_reward_signal(_grpo_args(), []) is False + + +# ---------------- the built-in filter is wired to the same question ---------------- + + +def test_builtin_filter_keeps_a_group_only_one_component_varies_in(): + from relax.engine.filters.dynamic_sampling_filters import check_reward_nonzero_std + + group = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[0.0, 1.0, 0.0, 1.0]) + assert check_reward_nonzero_std(_gdpo_args(), group).keep is True + + +def test_builtin_filter_drops_a_fully_flat_group(): + from relax.engine.filters.dynamic_sampling_filters import check_reward_nonzero_std + + group = _group(correctness=[1.0, 1.0, 1.0, 1.0], fmt=[0.5, 0.5, 0.5, 0.5]) + out = check_reward_nonzero_std(_gdpo_args(), group) + assert out.keep is False + assert out.reason.startswith("zero_std_") + + +def test_builtin_filter_is_unchanged_for_single_reward_algorithms(): + from relax.engine.filters.dynamic_sampling_filters import check_reward_nonzero_std + + flat = _group(correctness=[1.0, 1.0, 1.0, 1.0], fmt=[0.0, 0.0, 0.0, 0.0]) + assert check_reward_nonzero_std(_grpo_args(), flat).keep is False + + varied = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[1.0, 0.0, 1.0, 0.0]) + assert check_reward_nonzero_std(_grpo_args(), varied).keep is True From 884aa1a7e4a0a8664e010a0a0f1eab3cd5c9059d Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 16:49:49 +0800 Subject: [PATCH 05/22] fix(gdpo): correct the motivation, and the consumers that encoded it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A multi-CLI review found that the case this feature has been advertised on since the RFC is mathematically impossible, and that several guards were narrower than their own comments claimed. **The motivation was wrong.** "(correct, badly formatted) and (wrong, well formatted) both sum to 1, so GRPO flattens the group and GDPO does not" is false in its second half. If the components sum to a constant then `r2 = C - r1`, hence `std(r2) == std(r1)` and `z2 == -z1`, so equal weights cancel to exactly zero -- the same answer GRPO gives. No choice of samples changes this; only unequal weights break the tie. What GDPO actually buys, verified numerically and now documented: - *Relative strength between groups.* Per-group standardisation forces every group to unit variance, so a group where one component varies and a group where both vary come out identical under GRPO (both ±0.707). Standardising per component makes the second twice the amplitude, and step 3 whitens across the batch, so it survives (±0.548 vs ±1.095). - *Scale disparity between components.* correctness in {0,1} summed with a reward in the hundreds yields a sum whose variance is essentially the large component's. With the two ranked oppositely, GRPO gives the wrong-but-long responses the highest advantage; GDPO does not. **The consumers encoded the wrong question too.** `group_carries_reward_signal` returned True whenever any raw component varied, which disagrees with the advantage actually produced whenever a weight mutes a component or two components cancel -- it kept groups that then contribute no gradient. It now computes steps 1 and 2 (cheap, per-group) and asks whether the result is non-zero. **Single-reward judgement is back to float32.** The previous commit compared Python floats and described that as equivalent to the `std > 0` it replaced. It is not: `[0.1 + 0.2, 0.3, ...]` survives in float64 and collapses on cast, and a group of NaNs reads as varying. Both were kept and then contributed nothing. Judged on a float32 tensor, `min == max` and `std == 0` agree on every input, so this is now genuinely the old behaviour; non-finite groups are dropped as they were. **Overflow checks reached only the cast.** Individual rewards were bounded by float32's range, but the mean, std and weighted sum are computed in float32 too: [3e38, 2e38, 1e38, 0] overflows to -inf and gets silently zeroed. Weights had the same gap in both directions (1e300 casts to inf; [1e-50, 0] casts to all-zero after passing the "not all zero" check). All three now raise. **`_as_reward_tensor` leaked autograd.** `torch.tensor(x)` detached; `.to()` returns the same object when dtype and device match, so a reward tensor with `requires_grad=True` produced advantages carrying grad history. Detached. Also, from the same review: - `loss.py` chose the REINFORCE++ normalisation and mask-safe reducer by comparing algorithm names -- the last *maths* decision still made that way, and the one PR1 most conspicuously missed. Both now read a new `advantage_normalization` field. `loss.py` has no algorithm-name comparisons left. - `apply_custom_config_overrides` re-ran only the spec validator, so a YAML file could switch to `reinforce_plus_plus_baseline` and enable a reward hook it forbids: the spec is deliberately silent there, and the frozen validator never ran. It now runs both, in the main path's order. - With that fixed, `reinforce_plus_plus{,_baseline}` can declare the four constraints they were leaving at defaults. An undeclared field is not neutral -- it asserts the default -- so the spec had been stating four false things about them. - Docs claimed the built-in filter judges by `--reward-key` (untrue since the previous commit), and claimed exact scale invariance (the additive epsilon makes it approximate). The test oracle used a relative collapse tolerance while the implementation uses exact equality; binary rewards hid the disagreement. Tests: 1657 passed, same 2 failures + 2 errors as main@98a1274 here. --- docs/en/examples/algorithms.md | 12 +- docs/zh/examples/algorithms.md | 12 +- examples/gdpo/README.md | 25 +++- examples/gdpo/reward_gdpo.py | 35 +++-- examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh | 11 +- relax/algorithms/advantages.py | 29 ++++- relax/algorithms/policy.py | 3 +- relax/algorithms/rewards.py | 121 ++++++++++++++---- relax/algorithms/spec.py | 26 ++-- .../filters/dynamic_sampling_filters.py | 10 +- .../algorithms/test_arguments_spec_driven.py | 44 ++++++- tests/algorithms/test_gdpo.py | 11 +- .../algorithms/test_multi_reward_consumers.py | 101 +++++++++++++-- 13 files changed, 361 insertions(+), 79 deletions(-) diff --git a/docs/en/examples/algorithms.md b/docs/en/examples/algorithms.md index 3584482fc..3f6623523 100644 --- a/docs/en/examples/algorithms.md +++ b/docs/en/examples/algorithms.md @@ -299,7 +299,13 @@ The weights multiply the **normalized** advantages, not the raw rewards. After s $$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\mathrm{std}_\text{batch} + \epsilon}$$ -**Why this beats GRPO:** when one component is constant across a group (reward collapse), GRPO's summed reward collapses too, the whole group's advantages go to zero, and the samples are wasted. Under GDPO only *that component* contributes zero while the others still carry signal. +**Why this beats GRPO:** summing first and standardizing once discards two things. + +*Relative strength between groups.* Per-group standardization forces every group to unit variance, so a group where only one component varies and a group where both vary come out identical. Standardizing each component first makes the latter twice the amplitude, and step 3 whitens across the *batch*, so the difference survives into the final advantage. Measured (G=2, two groups in one batch): GRPO gives both groups ±0.707; GDPO gives ±0.548 and ±1.095. + +*Scale disparity between components.* A `correctness` in {0, 1} added to a reward in the hundreds (the paper's maths setup scores response length) yields a sum whose variance is essentially the large component's, so GRPO's direction is decided by it alone. GDPO gives each component unit variance first, so a weight expresses relative importance rather than units. + +**What GDPO does not do:** rescue a group whose components sum to a constant. There `r₂ = C − r₁` forces `z₂ = −z₁`, so equal weights cancel to exactly zero — the same answer GRPO gives. Only unequal weights break that tie. If *every* component is constant, GDPO returns zero as well. **On $\epsilon$:** GDPO uses $\epsilon = 10^{-4}$ at both steps, matching the reference implementation (the `scale_rewards` GDPO branch of TRL's `GRPOTrainer`), whereas GRPO / GSPO / SAPO / CISPO keep this repository's existing $10^{-6}$. The two only diverge on near-degenerate groups: with binary rewards and a group of 8 the within-group standard deviation is around 0.4 and the constants differ by 0.02%, but a continuous reward (the paper's maths setup scores response length) can leave a group at a standard deviation of ~$10^{-3}$, where $10^{-4}$ damps that group's signal by about 7% against 0.08% for $10^{-6}$. Groups that collapse *exactly* never reach this division; they are detected by exact equality and zeroed. @@ -348,7 +354,9 @@ Two differences between this implementation and the paper. Confirm they are acce - `--agentic-custom-advantage-path`: the second early return in `post_process_rewards`, which likewise returns ahead of the normalizer, with the same consequence. One flag, `AlgorithmSpec.allows_reward_post_process_hooks`, guards both. - `--fully-async`: see above. -All of these fail during argument validation. Combining it with `--dynamic-sampling-filter-path` logs a warning instead: the built-in `check_reward_nonzero_std` judges a group by the single `--reward-key` scalar and may drop groups whose signal lives in the other components. +All of these fail during argument validation. + +`--dynamic-sampling-filter-path` does **not** conflict: the built-in `check_reward_nonzero_std` is component-aware, computing what GDPO's first two steps actually produce and keeping the group only when that is non-zero, so its verdict matches the signal training receives. A warning is logged only for a *custom* filter, which may reduce the group to the single `--reward-key` scalar and drop groups whose signal lives in the other components. --- diff --git a/docs/zh/examples/algorithms.md b/docs/zh/examples/algorithms.md index 2e41a5260..5c5d26b57 100644 --- a/docs/zh/examples/algorithms.md +++ b/docs/zh/examples/algorithms.md @@ -296,7 +296,13 @@ $$A_\text{sum}^{(i,j)} = \sum_k w_k A_k^{(i,j)}$$ $$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\mathrm{std}_\text{batch} + \epsilon}$$ -**相对 GRPO 的收益**：当一组 rollout 的各**分量不同、但总和恰好相同**时（例如 `(1,0)` 与 `(0,1)` 都求和为 1），GRPO 看到的组内总奖励无差异、整组 advantage 归零被丢弃；GDPO 对每个分量单独标准化，仍能保留各分量的学习信号。若某个分量在组内恒定，则只有**该分量**贡献 0、其它分量照常提供信号；若**所有**分量都恒定，GDPO 与 GRPO 一样返回零。 +**相对 GRPO 的收益**：GRPO 把各分量相加后只做一次组内标准化，这会丢掉两类信息。 + +*一是组间的相对强度*：组内标准化把每一组都拉到单位方差，于是「只有一个分量在变」的组与「两个分量都在变」的组得到完全相同的 advantage。GDPO 让各分量先各自标准化，后者幅度自然是两倍；而第三步的 batch 白化是**跨组**的，这个差异会保留到最终 advantage。实测（G=2，一个 batch 两组）：GRPO 两组同为 ±0.707，GDPO 分别为 ±0.548 与 ±1.095。 + +*二是分量间的尺度差异*：`correctness ∈ {0,1}` 与一个取值上百的分量（论文实验用响应长度）相加时，和的方差几乎全部来自后者，GRPO 的方向由它单独决定；GDPO 先让每个分量单位方差，权重才真正表达相对重要性而非量纲。 + +**GDPO 做不到什么**：若各分量在组内**恰好加和为常数**（`r₂ = C − r₁`），标准化后恒有 `z₂ = −z₁`，等权重下**完全抵消为零**，与 GRPO 结果相同——只有不等权重能在这类组上取得信号。若**所有**分量都恒定，GDPO 同样返回零。 **关于 $\epsilon$**：GDPO 的两步都用 $\epsilon = 10^{-4}$，与参考实现（TRL `GRPOTrainer` 的 `scale_rewards` GDPO 分支）一致，而 GRPO / GSPO / SAPO / CISPO 沿用本仓库既有的 $10^{-6}$。两者的差别只在近乎塌缩的组上显现：二值 reward、组大小 8 时组内标准差约 0.4，两个取值的差异是 0.02%；但连续 reward（论文的数学实验用响应长度）可能让某组的标准差落到 $10^{-3}$ 量级，此时 $10^{-4}$ 会把该组的信号额外压低约 7%，而 $10^{-6}$ 只压低 0.08%。**完全**塌缩的组不会走到这个除法——它们由 exact 相等判定后直接置零。 @@ -345,7 +351,9 @@ GDPO_ARGS=( - 不能与 `--agentic-custom-advantage-path` 同用：`post_process_rewards` 里的第二个早返回点，同样赶在归一化器之前返回，后果与上一条相同。两者由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守。 - 不能与 `--fully-async` 同用（见上）。 -以上都会在参数校验阶段直接报错。配合 `--dynamic-sampling-filter-path` 时会给出警告：内置的 `check_reward_nonzero_std` 只看 `--reward-key` 那一个标量，可能丢掉只存在于其它分量的信号。 +以上都会在参数校验阶段直接报错。 + +`--dynamic-sampling-filter-path` **不**冲突：内置的 `check_reward_nonzero_std` 已经是分量感知的，它直接算出 GDPO 前两步的组合结果、按其是否非零判定，因此与训练实际拿到的信号一致。只有指向**自定义** filter 时才会给出警告——那种 filter 若只看 `--reward-key` 标量，就会丢掉只存在于其它分量的信号。 --- diff --git a/examples/gdpo/README.md b/examples/gdpo/README.md index 7216e3a48..676bfa941 100644 --- a/examples/gdpo/README.md +++ b/examples/gdpo/README.md @@ -11,9 +11,24 @@ 这两个分量会**不同步**：模型可能答对但没按格式输出，也可能格式完美但答错。 -GRPO 把它们**加起来**再做一次组内归一化。问题是不同的分量组合可能得到**相同的总和**：一组 rollout 里，有的是「答对但格式差」`(correctness=1, format=0)`、有的是「答错但格式好」`(0, 1)`，总奖励都等于 1——GRPO 看到组内总奖励全相同，整组 advantage 归零、样本白采，可两个分量各自明明都有信号。 +GRPO 把它们**加起来**再做一次组内归一化。这一步会丢掉两类信息。 -GDPO 分别对每个分量做组内标准化再合并，就能区分这两类样本，`correctness` 与 `format` 各自的差异都会转成梯度信号。（反过来，若两个分量在组内**都**恒定，GDPO 与 GRPO 一样返回零，不会无中生有。） +**一、组间的相对强度。** 组内标准化把每一组的 advantage 都拉到单位方差，于是「只有 correctness 在变」的组和「两个分量都在变」的组，得到完全相同的 advantage。GDPO 对每个分量各自标准化后相加，两个分量都起作用时幅度自然是两倍；而第三步的 batch 白化是**跨组**的，所以这个差异会一路保留到最终 advantage： + +| | 组 A（只有 correctness 变化） | 组 B（两个分量都变化） | +| ---------------- | ----------------------------- | ---------------------- | +| GRPO | ±0.707 | ±0.707（分不出） | +| GDPO（含第三步） | ±0.548 | ±1.095（B 强一倍） | + +**二、分量之间的尺度差异。** `correctness ∈ {0,1}` 与一个取值上百的 `format`（或论文实验里的响应长度）相加时，和的方差几乎全部来自大尺度那一维，GRPO 的方向就由它单独决定。GDPO 先让每个分量单位方差，权重才真正表达「相对重要性」而不是量纲。极端一点：`correctness=[1,1,0,0]`、`format=[0,100,200,300]`（两者排序相反）时，GRPO 给答错的长响应最高 advantage，GDPO 不会。 + +## GDPO 不能做什么 + +**各分量在组内恰好加和为常数时，GDPO 也救不回来。** 若 `correctness + format ≡ C`，则 `format = C − correctness`，两者标准化后恒有 `z_format = −z_correctness`，**等权重下完全抵消为零**——与 GRPO 得到同样的结果。 + +这一点本文档此前写反了，说这正是 GDPO 的优势场景。它不是：那是一个数学上不可能被等权重 GDPO 区分的情形。只有不等权重（如 `--gdpo-reward-weights 2.0 1.0`）能在这类组上拿到信号。 + +（另外，若两个分量在组内**都**恒定，GDPO 与 GRPO 一样返回零，不会无中生有。） ## 运行 @@ -53,7 +68,9 @@ GDPO_ARGS=( ) ``` -`--gdpo-reward-weights` 乘的是**归一化之后**的 advantage，不是原始 reward。经过第一步各分量已经是单位方差，所以权重表达的是相对重要性，与分量本身的量纲无关——把 `format` 的取值范围从 `[0,1]` 改成 `[0,100]` 不会改变训练结果。 +`--gdpo-reward-weights` 乘的是**归一化之后**的 advantage，不是原始 reward。经过第一步各分量已经是单位方差，所以权重表达的是相对重要性，而不是分量的量纲——把 `format` 的取值范围从 `[0,1]` 改成 `[0,100]` 基本不改变训练结果。 + +说「基本」而不是「完全」，是因为第一步除的是 `std + 1e-4` 而不是 `std`。当某分量的组内标准差本身就落到 `1e-4` 量级时，这个加性 epsilon 会随量纲变化而改变阻尼比例，缩放就不再是严格等价的。二值奖励（std≈0.4）离这个区间很远，连续奖励（如响应长度）则可能撞上。 ## 换成自己的奖励 @@ -71,4 +88,4 @@ GDPO_ARGS=( `--normalize-advantages`、`--custom-reward-post-process-path`、`--agentic-custom-advantage-path` 和 `--fully-async` 都不能与 GDPO 同用，参数校验阶段会直接报错。第一个会造成双重白化；第二、三个都会赶在归一化器之前从 `post_process_rewards` 返回，导致 GDPO 的前两步被静默跳过（这两个由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守）；第四个的统计窗口可能小到只剩一个样本。 -配 `--dynamic-sampling-filter-path` 时只警告不报错：内置过滤器按 `--reward-key` 的单个标量判组，可能丢掉只在其它分量里有信号的组。 +配 `--dynamic-sampling-filter-path` 时**不冲突**：内置的 `check_reward_nonzero_std` 会直接算出 GDPO 前两步的组合结果、按其是否非零判组，与训练实际拿到的信号一致（权重把某分量静音、或两个分量抵消，它都能算出来）。只有指向**自定义** filter 时才会警告——那种 filter 若只看 `--reward-key` 标量，就会丢掉只在其它分量里有信号的组。 diff --git a/examples/gdpo/reward_gdpo.py b/examples/gdpo/reward_gdpo.py index 476d6cee0..f0e694792 100644 --- a/examples/gdpo/reward_gdpo.py +++ b/examples/gdpo/reward_gdpo.py @@ -2,13 +2,24 @@ """Two-component reward for the GDPO example: correctness and format. -GDPO standardizes each component within its prompt group before combining them, -so a group whose rollouts differ in their reward *components* but share the same -*summed* reward still carries a learning signal. Example: (correct, badly -formatted) and (wrong, well formatted) both sum to 1 -- GRPO sees one constant -summed reward and the whole group contributes nothing, while GDPO keeps each -component's signal. (If every component is constant within the group, GDPO -returns zero too.) +GDPO standardizes each component within its prompt group before combining them. +Two things survive that GRPO's single standardization of the summed reward +destroys: + +* **Relative strength between groups.** Standardizing per group forces every + group to unit variance, so a group where only `correctness` varies and one + where both components vary come out identical. Standardizing each component + first makes the latter twice the amplitude, and step 3 whitens across the + *batch*, so the difference reaches the final advantage. +* **Scale disparity between components.** A `correctness` in {0, 1} added to a + reward in the hundreds gives a sum whose variance is essentially the large + component's, so GRPO's direction is decided by it alone. + +What GDPO does *not* do: rescue a group whose components sum to a constant. +There `format = C - correctness` forces the standardized values to be exact +opposites, and equal weights cancel them to zero -- the same answer GRPO +gives. Only unequal weights break that tie. (If every component is constant +within the group, GDPO returns zero too.) Wire it up with:: @@ -55,8 +66,14 @@ def compute_gdpo_reward(response: str, label: Any) -> dict[str, float]: """Score one response on answer correctness and on output format. The two components are deliberately decorrelated: a response can be correct - without the expected tags, and well-formatted while wrong. That is the - situation GDPO handles better than a summed reward. + without the expected tags, and well-formatted while wrong. + + ``score`` is ``correctness`` rather than the sum, on purpose. It feeds + ``--reward-key``, which selects the scalar for metrics and the + ``raw_reward`` column only -- it does not participate in the GDPO + computation, which reads the two components directly. Reporting accuracy + there is more legible on a dashboard than a blended number. Note this means + ``rollout/raw_reward`` tracks correctness alone, not overall reward. """ answer = _extract_answer(response) correctness = 1.0 if answer is not None and answer == _final_answer(label) else 0.0 diff --git a/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh index 447d2f555..5c9a04442 100644 --- a/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh +++ b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh @@ -5,10 +5,13 @@ # Qwen3-0.6B single-GPU GDPO training on GSM8K. # # GDPO (arXiv 2601.05242) standardizes each reward component within its prompt -# group before combining them, so a group whose rollouts share the same summed -# reward but differ in their components still carries signal -- e.g. (correct, -# badly formatted) and (wrong, well formatted) both sum to 1, which GRPO flattens -# to nothing. The reward function in reward_gdpo.py returns both components; +# group before combining them. That keeps two things GRPO's single +# standardization of the summed reward loses: the relative strength between +# groups (per-group standardization forces every group to unit variance), and +# the balance between components of very different scale. It does NOT rescue a +# group whose components sum to a constant -- those cancel to zero under equal +# weights, same as GRPO. See examples/gdpo/README.md. +# The reward function in reward_gdpo.py returns both components; # --gdpo-reward-keys names them. # # Usage: diff --git a/relax/algorithms/advantages.py b/relax/algorithms/advantages.py index 99fd755be..2c6936c1c 100644 --- a/relax/algorithms/advantages.py +++ b/relax/algorithms/advantages.py @@ -66,8 +66,18 @@ def whiten_scalar(values: torch.Tensor, *, process_group: dist.ProcessGroup | No def _as_reward_tensor(rewards: Any, kl: list[torch.Tensor]) -> torch.Tensor: + """Rewards as a detached float32 tensor on the KL tensors' device. + + ``detach()`` is load-bearing, not defensive. The pre-registry code built + this with ``torch.tensor(rewards, ...)``, which copies and drops autograd + history even when handed a tensor. ``.to()`` returns the *same* object when + dtype and device already match, so without the detach a caller passing a + reward tensor with ``requires_grad=True`` would get advantages that carry + grad history into the policy loss -- a path that did not exist before and + that no caller asks for. Rewards are data, not something to backprop into. + """ if isinstance(rewards, torch.Tensor): - return rewards.to(dtype=torch.float32, device=kl[0].device) + return rewards.detach().to(dtype=torch.float32, device=kl[0].device) return torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) @@ -137,6 +147,23 @@ def _whiten_by_segment(values, mini_batch_sizes, process_group): data, so every rank is expected to run the same number of segments; :func:`_agree_on_segmentation` is what turns that expectation into a checked precondition rather than a deadlock. + + What it checks is the segment *count*, not segment *identity*. Two further + properties are relied on and not verified here: + + * Segment ``k`` holds training batch ``k`` on every rank. ``actor.py`` + fetches batches in ``batch_index`` order and appends the counts in that + same order, so the orders agree; nothing in this function would notice if + they stopped agreeing, and the failure would be silent -- statistics + mixed across two training batches, no error. + * A segment's sample count may differ between ranks (a data-parallel split + balances tokens, not samples) and that is fine, because the statistic is + reduced across the group. What is not fine is two ranks disagreeing about + *which* batch a segment belongs to. + + Both are guaranteed upstream rather than here because checking identity + would need a batch id in the metadata, which the plan does not currently + carry. If that metadata ever appears, fold it into the same MAX reduction. """ # Validate whatever was passed, including a single segment: treating an empty # or malformed list as "fall back to one window" would silently restore the diff --git a/relax/algorithms/policy.py b/relax/algorithms/policy.py index 12e023bbc..dc2974123 100644 --- a/relax/algorithms/policy.py +++ b/relax/algorithms/policy.py @@ -23,7 +23,8 @@ def policy_loss_ppo_clip(args: Any, *, log_probs, ppo_kl, advantages): - """Standard clipped surrogate objective (GRPO, GSPO, GDPO, REINFORCE++).""" + """Standard clipped surrogate objective (GRPO, GSPO, GDPO, PPO, + REINFORCE++).""" return compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high) diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index f04d00855..59ce5f3d5 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -222,6 +222,20 @@ def extract_reward_components(samples: list[Any], keys: list[str]) -> torch.Tens return components +def standardize_group_components(group: torch.Tensor) -> torch.Tensor: + """GDPO step 1 on one ``[G, K]`` prompt group (arXiv 2601.05242, Eq. 4). + + Each column is standardised on its own. A column that is flat across the + group contributes exactly zero rather than ``0 / (0 + eps)`` noise, which + is what lets the remaining columns keep their signal. + """ + centered = group - group.mean(dim=0, keepdim=True) + std = group.std(dim=0) + collapsed = collapsed_columns(group, dim=0) + scaled = centered / (std + GDPO_EPS) + return torch.where(collapsed.unsqueeze(0), torch.zeros_like(scaled), scaled) + + def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: """GDPO steps 1 and 2 (arXiv 2601.05242, Eq. 4 and Eq. 7). @@ -231,10 +245,20 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl TransferQueue schema change. Step 3 (batch whitening) runs later, in :func:`relax.algorithms.advantages.advantage_gdpo`. - Standardising per component before combining is what separates GDPO from - GRPO: when one component collapses within a group, only that component - contributes zero, while GRPO's summed reward would collapse and discard the - whole group. + What standardising per component actually buys, stated carefully because + it is easy to overclaim: the combined advantage is ``sum_k w_k * z_k``, + where each ``z_k`` has unit variance within the group. GRPO instead + standardises ``sum_k r_k``, so a component with a large spread dominates + the direction. When the components have comparable spread the two agree up + to a positive scalar; they diverge when the spreads differ, which is the + case GDPO is for -- a correctness reward in ``{0, 1}`` combined with a + length reward in the hundreds is decided almost entirely by length under + GRPO, and half by each under GDPO. + + It does **not** rescue a group whose components sum to a constant. There + ``r_2 = C - r_1`` forces ``z_2 = -z_1``, so equal weights cancel to exactly + zero -- the same answer GRPO gives. Unequal weights break the tie; equal + weights cannot. """ keys = resolve_gdpo_keys(args) weights = resolve_gdpo_weights(args, keys) @@ -246,12 +270,8 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl fully_collapsed_groups = 0 for positions in positions_by_group.values(): group = components[positions] - centered = group - group.mean(dim=0, keepdim=True) - std = group.std(dim=0) - collapsed = collapsed_columns(group, dim=0) - scaled = centered / (std + GDPO_EPS) - normalized[positions] = torch.where(collapsed.unsqueeze(0), torch.zeros_like(scaled), scaled) - if bool(collapsed.all()): + normalized[positions] = standardize_group_components(group) + if bool(collapsed_columns(group, dim=0).all()): fully_collapsed_groups += 1 if fully_collapsed_groups == len(positions_by_group): @@ -269,7 +289,36 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl ) weight_tensor = torch.tensor(weights, dtype=torch.float32) - return (normalized * weight_tensor).sum(dim=1).tolist() + if not torch.isfinite(weight_tensor).all(): + # `resolve_gdpo_weights` checked `math.isfinite` on the Python floats. + # That is a different question: 1e300 is finite in float64 and `inf` + # once it lands in this tensor. + raise ValueError( + f"--gdpo-reward-weights {weights} contains a value that overflows float32; " + f"the largest representable is {_FLOAT32_MAX:.6g}." + ) + if float(weight_tensor.abs().sum()) == 0.0: + # Same gap in the other direction: 1e-50 is a nonzero Python float that + # flushes to zero in float32, so a weight vector that passed the + # "not all zero" check can still be all zeros by the time it multiplies. + raise ValueError( + f"--gdpo-reward-weights {weights} are all zero once cast to float32; " + "the combined advantage would be identically 0." + ) + + combined = (normalized * weight_tensor).sum(dim=1) + if not torch.isfinite(combined).all(): + # Every individual reward fit in float32 (checked in + # `extract_reward_components`), but the arithmetic between them need + # not: a group of [3e38, 2e38, 1e38, 0] overflows while computing its + # own mean. Left unchecked the `inf` reaches `whiten_scalar`, reads as + # a non-finite std, and the batch is silently zeroed -- the exact + # failure the per-value check was added to prevent, one stage later. + raise ValueError( + "GDPO produced a non-finite combined advantage from finite inputs; the float32 " + f"arithmetic overflowed. Rescale the rewards (keys={keys}) or their weights." + ) + return combined.tolist() def group_carries_reward_signal(args: Any, samples: list[Any]) -> bool: @@ -281,30 +330,50 @@ def group_carries_reward_signal(args: Any, samples: list[Any]) -> bool: single scalar ``--reward-key`` selects, which is the right question only for an algorithm that consumes that scalar. - A multi-reward algorithm standardises each component separately, so a group - is dead only when *every* component is flat. Judging GDPO by the summed - scalar throws away exactly the groups it exists to keep: ``(1, 0)`` and - ``(0, 1)`` have identical sums and different components. - - The two consumers did not previously agree on the single-reward test: - ``check_reward_nonzero_std`` used ``std > 0``, the zero-std metrics used - exact equality. This unifies them on exact equality, matching - :func:`relax.algorithms.numerics.is_collapsed` and for the same reason -- - a tolerance wide enough to absorb float error also discards real signal, - and the two answers differ only for a group that is exactly flat yet whose - computed standard deviation is not exactly 0. + For a multi-reward algorithm the honest test is not "did any component + vary" but "is the combined advantage this algorithm will actually compute + non-zero". Those differ: a zero weight mutes a varying component, and two + components whose standardised values are exact opposites cancel. Asking + the cheaper question keeps groups that then contribute no gradient, and + under-reports the zero-std metrics. Steps 1 and 2 are per-group and cheap, + so this runs them. + + The single-reward test is deliberately performed **in float32**, on a + tensor, rather than on the Python floats. That is the precision the reward + actually reaches training in (:func:`_group_normalize` casts to float32, + and so does the GDPO path), so it is the precision that decides whether a + group will still carry signal by the time it matters. + + Comparing the float64 values instead is not a harmless tightening: it flips + the verdict for any group whose spread survives in float64 but collapses on + cast -- ``[0.1 + 0.2, 0.3, ...]`` is the everyday example, and a group of + NaNs is the pathological one (``nan != nan`` reads as signal). Both would + be kept and then contribute a zero gradient. Judging on the tensor keeps + this identical to the ``std > 0`` test it replaces: ``min == max`` and + ``std == 0`` agree on every float32 input. """ if not samples: return False spec = get_algorithm(args.advantage_estimator) if not spec.uses_reward_components: - rewards = [sample.get_reward_value(args) for sample in samples] - return any(reward != rewards[0] for reward in rewards) + rewards = torch.tensor([sample.get_reward_value(args) for sample in samples], dtype=torch.float32) + if not torch.isfinite(rewards).all(): + # A group containing NaN or inf reads as "varying" under any + # inequality test (`nan != nan` is True), which would forward a + # broken reward into training. `std > 0` answered False here + # because the std is itself NaN, so False preserves that. Raising + # would arguably be better -- a non-finite reward is a bug in the + # reward function, not a property of the group -- but that is a + # behaviour change this refactor is not the place for. + return False + return bool(rewards.amin() != rewards.amax()) keys = resolve_gdpo_keys(args) + weights = torch.tensor(resolve_gdpo_weights(args, keys), dtype=torch.float32) components = extract_reward_components(samples, keys) - return bool((~collapsed_columns(components, dim=0)).any()) + combined = (standardize_group_components(components) * weights).sum(dim=1) + return bool((combined != 0).any()) REWARD_NORMALIZERS: dict[str, Callable[[Any, list[Any], list[float]], list[float]]] = { diff --git a/relax/algorithms/spec.py b/relax/algorithms/spec.py index 7c76d9285..446a214a7 100644 --- a/relax/algorithms/spec.py +++ b/relax/algorithms/spec.py @@ -259,6 +259,13 @@ def is_group_normalized(self) -> bool: advantage_fn="reinforce_plus_plus", policy_loss_fn="ppo_clip", requires_normalize_advantages=True, + # `_validate_reinforce_plus_plus_args` also rejects --fully-async for + # this estimator. Declaring it here as well is not redundant: an + # undeclared field defaults to True, i.e. the spec would actively state + # something false about the algorithm, and any future reader of the + # spec would believe it. That function still runs first and still owns + # the wording its tests match on. + supports_fully_async=False, ), "reinforce_plus_plus_baseline": AlgorithmSpec( name="reinforce_plus_plus_baseline", @@ -267,15 +274,18 @@ def is_group_normalized(self) -> bool: 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. + # These five are also enforced by hand in + # `_validate_reinforce_plus_plus_args`, a frozen Task 29 contract that + # owns the wording its tests match on. They were left undeclared at + # first to avoid replacing that wording -- but an undeclared field is + # not neutral, it *asserts the default*, so the spec was stating five + # things about this estimator that are false. Both validation paths now + # run the frozen function first (the main one always did; + # `apply_custom_config_overrides` was fixed to match), so its messages + # still win and the spec can stop lying. + supports_fully_async=False, requires_rewards_normalization=True, + allows_reward_post_process_hooks=False, min_group_size=2, forbids_reward_side_kl=True, ), diff --git a/relax/engine/filters/dynamic_sampling_filters.py b/relax/engine/filters/dynamic_sampling_filters.py index 501e352df..b2425b359 100644 --- a/relax/engine/filters/dynamic_sampling_filters.py +++ b/relax/engine/filters/dynamic_sampling_filters.py @@ -12,10 +12,12 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): """Drop prompt groups whose rewards carry no signal for this algorithm. "No signal" is algorithm-dependent, which is why the test is not spelled - out here. For a single-reward algorithm it is the standard deviation of the - ``--reward-key`` scalar, exactly as before. For a multi-reward algorithm it - is whether *every* component is flat -- judging those by the summed scalar - would drop the groups the algorithm exists to keep. + out here. For a single-reward algorithm it is whether the ``--reward-key`` + scalar varies in float32 -- equivalent to the ``std > 0`` this replaced, + since ``min == max`` and ``std == 0`` agree on every float32 input. For a + multi-reward algorithm it is whether *every* component is flat; judging + those by the summed scalar would drop the groups the algorithm exists to + keep. """ keep = group_carries_reward_signal(args, samples) return DynamicFilterOutput( diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py index 3764ffb82..f1a2944ca 100644 --- a/tests/algorithms/test_arguments_spec_driven.py +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -441,11 +441,18 @@ def test_other_estimators_are_unaffected_by_fully_async(arguments_module, estima arguments_module.validate_algorithm_args(args) -def test_supports_fully_async_defaults_to_true(): +def test_supports_fully_async_is_declared_wherever_it_is_false(): + """The set must match what validation actually rejects. + + An undeclared field defaults to True, which is an assertion, not silence: + the REINFORCE++ variants are rejected under --fully-async by + `_validate_reinforce_plus_plus_args`, so leaving them at the default would + make the spec state the opposite of the behaviour. + """ from relax.algorithms import get_algorithm, list_algorithm_names unsupported = {n for n in list_algorithm_names() if not get_algorithm(n).supports_fully_async} - assert unsupported == {"gdpo"} + assert unsupported == {"gdpo", "reinforce_plus_plus", "reinforce_plus_plus_baseline"} def test_dynamic_sampling_filter_warns_for_multi_reward(arguments_module, caplog): @@ -465,3 +472,36 @@ def test_no_warning_without_a_filter(arguments_module, caplog): arguments_module.validate_algorithm_args(_args()) assert not any("dynamic-sampling-filter-path" in r.message for r in caplog.records) + + +def test_yaml_cannot_bypass_the_frozen_reinforce_plus_plus_contract(arguments_module, tmp_path): + """The YAML re-check must run *both* validators, like the main path does. + + `reinforce_plus_plus_baseline` keeps three of its constraints in + `_validate_reinforce_plus_plus_args` instead of in its spec, precisely so + that function keeps owning its frozen wording. That decision opens a hole + if the override path re-runs only the spec-driven validator: a YAML file + could switch to that estimator and enable a reward hook it forbids, and + neither validator would object -- the spec because it is deliberately + silent, the frozen one because it never ran. + """ + args = _overridable_args( + tmp_path, + "advantage_estimator: reinforce_plus_plus_baseline\ncustom_reward_post_process_path: foo.py\n", + normalize_advantages=True, + use_kl_loss=True, + kl_loss_coef=0.01, + kl_loss_type="k2", + kl_coef=0.0, + n_samples_per_prompt=4, + use_unbiased_kl=False, + # Everything else that validator insists on, so the assertion below is + # about the reward hook and not about whichever check fires first. + colocate=True, + fully_async=False, + hybrid=False, + context_parallel_size=1, + calculate_per_token_loss=False, + ) + with pytest.raises(ValueError, match="custom-reward-post-process-path"): + arguments_module.apply_custom_config_overrides(args) diff --git a/tests/algorithms/test_gdpo.py b/tests/algorithms/test_gdpo.py index e72884af8..de5dbea2a 100644 --- a/tests/algorithms/test_gdpo.py +++ b/tests/algorithms/test_gdpo.py @@ -72,8 +72,15 @@ def _manual_gdpo(correctness, fmt, groups, weights=(1.0, 1.0)): mean = sum(vals) / len(vals) var = sum((v - mean) ** 2 for v in vals) / (len(vals) - 1) std = math.sqrt(var) - scale = max(abs(v) for v in vals) - collapsed = std <= 1e-6 * scale + # Exact equality, matching the implementation. A relative tolerance + # here would make the oracle disagree with the code on precisely the + # inputs where the choice matters -- a narrow continuous column with + # std near 1e-6 * scale, which the oracle would zero and the + # implementation would keep. The existing cases are binary rewards, + # whose std is either exactly 0 or >= 0.5, so the two criteria never + # diverge there and the mismatch would go unnoticed until someone + # added a continuous case and mistrusted the wrong side. + collapsed = max(vals) == min(vals) # 1e-4, hardcoded on purpose: this oracle exists to pin parity with # the reference implementation (trl grpo_trainer.py's scale_rewards # GDPO path divides by std + 1e-4 at both steps), so importing our diff --git a/tests/algorithms/test_multi_reward_consumers.py b/tests/algorithms/test_multi_reward_consumers.py index 3ab588531..2a2badccd 100644 --- a/tests/algorithms/test_multi_reward_consumers.py +++ b/tests/algorithms/test_multi_reward_consumers.py @@ -4,14 +4,20 @@ The dynamic-sampling filter and the zero-std metrics both decide whether a prompt group "carries signal". They answered that with the single -``--reward-key`` scalar, which is only the right question for an algorithm that +``--reward-key`` scalar, which is the right question only for an algorithm that consumes that scalar. -For a multi-reward algorithm it is the wrong one, and wrong in the direction -that undoes the algorithm: a group where ``correctness`` is flat but ``format`` -still varies looks dead by the summed scalar, gets dropped by the filter, and -never reaches training -- while GDPO would have extracted signal from the -component that did vary. Keeping those groups is the entire point. +For a multi-reward algorithm the right question is narrower than "did any +component vary" and wider than "did the scalar vary": it is whether the +*combined* advantage GDPO will actually produce is non-zero. The three differ. +A zero weight mutes a varying component; two components whose standardised +values are exact opposites cancel; and a group can be flat in the scalar while +alive in the components. Only the combined value predicts whether the group +contributes a gradient, so that is what these consumers compute. + +The tests below deliberately include the case that motivated all of this and +turned out to be wrong -- equal sums under equal weights cancel to zero -- so +that the mistake cannot be reintroduced as a "fix". """ from types import SimpleNamespace @@ -21,7 +27,12 @@ torch = pytest.importorskip("torch") -from relax.algorithms.rewards import group_carries_reward_signal # noqa: E402 +from relax.algorithms.rewards import group_carries_reward_signal, normalize_gdpo_decoupled # noqa: E402 + + +def _combined(args, group): + """The scalar GDPO steps 1+2 actually hand to the advantage stage.""" + return normalize_gdpo_decoupled(args, group, [0.0] * len(group)) class _S: @@ -71,17 +82,55 @@ def test_group_flat_in_one_component_still_carries_signal(): assert group_carries_reward_signal(_gdpo_args(), group) is True -def test_equal_sums_from_different_components_carry_signal(): - """(1,0) and (0,1) both sum to 1 -- the paper's motivating case. +def test_equal_sums_cancel_exactly_under_equal_weights(): + """Equal sums do NOT survive GDPO under equal weights. Pinning the maths. - Every sample's --reward-key scalar is identical, so this is precisely the - group a scalar-based filter throws away. + This was written the other way round first, and the mistake is worth a + test of its own: if the components sum to a constant then ``r2 = C - r1``, + so ``std(r2) == std(r1)`` and ``z2 == -z1``. Equal weights cancel them to + exactly zero -- the same answer GRPO gives. Nothing about GDPO rescues + this group, and a filter that claims otherwise keeps a group that then + contributes no gradient. """ group = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[0.0, 1.0, 0.0, 1.0]) scalars = {s.get_reward_value(_gdpo_args()) for s in group} - assert scalars == {1.0}, "precondition: the scalar really is constant" + assert scalars == {1.0}, "precondition: the summed scalar really is constant" - assert group_carries_reward_signal(_gdpo_args(), group) is True + combined = _combined(_gdpo_args(), group) + assert combined == [0.0, 0.0, 0.0, 0.0], combined + assert group_carries_reward_signal(_gdpo_args(), group) is False + + +def test_unequal_weights_break_the_cancellation(): + """The tie is broken by weights, not by the standardisation itself.""" + args = _gdpo_args() + args.gdpo_reward_weights = [2.0, 1.0] + group = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[0.0, 1.0, 0.0, 1.0]) + + assert any(v != 0.0 for v in _combined(args, group)) + assert group_carries_reward_signal(args, group) is True + + +def test_scale_disparity_is_what_separates_gdpo_from_grpo(): + """The real motivation: components whose spreads differ by orders of + magnitude. + + `correctness` in {0,1} against a `format` reward in the hundreds, ranked + the opposite way. GRPO standardises the *sum*, so the large-spread + component decides the direction and even reverses the sign for the correct + samples. GDPO gives each component unit variance first, so both get a say. + """ + args = _gdpo_args() + group = _group(correctness=[1.0, 1.0, 0.0, 0.0], fmt=[0.0, 100.0, 200.0, 300.0]) + + summed = torch.tensor([s.get_reward_value(args) for s in group]) + grpo = (summed - summed.mean()) / (summed.std() + 1e-6) + gdpo = torch.tensor(_combined(args, group)) + + # GRPO ranks the two correct samples *lowest*, dragged there by `format`. + assert grpo[0] < grpo[2] and grpo[1] < grpo[3] + # GDPO does not: correctness pulls them back up, and the signs disagree. + assert not torch.equal(grpo.sign(), gdpo.sign()) def test_group_flat_in_every_component_is_dead(): @@ -111,12 +160,36 @@ def test_empty_group_carries_nothing(): def test_builtin_filter_keeps_a_group_only_one_component_varies_in(): + """`correctness` flat, `format` varying: one live component is enough.""" from relax.engine.filters.dynamic_sampling_filters import check_reward_nonzero_std - group = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[0.0, 1.0, 0.0, 1.0]) + group = _group(correctness=[1.0, 1.0, 1.0, 1.0], fmt=[0.0, 0.5, 1.0, 0.5]) assert check_reward_nonzero_std(_gdpo_args(), group).keep is True +def test_builtin_filter_agrees_with_the_advantage_it_will_produce(): + """The filter's verdict must match what the reward stage actually outputs. + + Cheaper proxies ("did any raw component vary?") disagree with the real + answer whenever a weight mutes a component or two standardised components + cancel -- and disagreeing means keeping groups with no gradient. + """ + from relax.engine.filters.dynamic_sampling_filters import check_reward_nonzero_std + + muted = _gdpo_args() + muted.gdpo_reward_weights = [0.0, 1.0] + cases = [ + (_gdpo_args(), _group([1.0, 0.0, 1.0, 0.0], [0.0, 1.0, 0.0, 1.0])), # cancels + (_gdpo_args(), _group([1.0, 1.0, 0.0, 0.0], [0.0, 100.0, 200.0, 300.0])), # real signal + (_gdpo_args(), _group([1.0, 1.0, 1.0, 1.0], [0.5, 0.5, 0.5, 0.5])), # fully flat + (muted, _group([1.0, 0.0, 1.0, 0.0], [2.0, 2.0, 2.0, 2.0])), # only muted one varies + ] + for args, group in cases: + verdict = check_reward_nonzero_std(args, group).keep + actually_moves = any(v != 0.0 for v in _combined(args, group)) + assert verdict is actually_moves, (verdict, _combined(args, group)) + + def test_builtin_filter_drops_a_fully_flat_group(): from relax.engine.filters.dynamic_sampling_filters import check_reward_nonzero_std From afe13f5d870f8bff7dd59c12236391bba5ac7b71 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 17:09:22 +0800 Subject: [PATCH 06/22] test(algorithms): make two assertions capable of failing Both were flagged in review as tautologies, and both were. `test_both_paths_delegate_to_the_shared_estimator` asserted only that the import line exists. A file that imports the shared handler and then ignores it satisfies that -- which is what a botched revert looks like. It now asserts the call, plus the absence of any direct `ppo_utils` kernel in either path: if either grows its own `get_grpo_returns` again, the estimator is being chosen outside the registry. Verified by mutation: re-adding that import to loss.py turns the test red. `test_no_estimator_name_comparisons_remain` claimed nothing in arguments.py compares algorithm names while checking two literals -- and `_validate_reinforce_plus_plus_args` does compare names, deliberately, as a frozen contract. Widening it to the whole file would either fail or need an exception list that hides the next regression. It is now scoped to `validate_algorithm_args`, named accordingly, and reads the function through `ast.unparse` so the check sees code rather than the docstring's own description of the pattern it removed. Any comparison inside that function fails it now, not just two spellings. Tests: 1657 passed, same 2 failures + 2 errors as main@98a1274 here. --- .../algorithms/test_arguments_spec_driven.py | 34 +++++++++++++++---- tests/algorithms/test_policy_loss_dispatch.py | 18 +++++++++- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py index f1a2944ca..2df3aded1 100644 --- a/tests/algorithms/test_arguments_spec_driven.py +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -79,13 +79,33 @@ def test_no_hardcoded_estimator_choice_list(): 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_spec_driven_validation_makes_no_name_comparisons(): + """Scoped to `validate_algorithm_args`, and the name says so. + + The previous version of this test claimed no name comparison remained + anywhere in arguments.py while only checking two literals -- and + `_validate_reinforce_plus_plus_args` does compare names, deliberately, as a + frozen contract. Asserting over the whole file would either fail or require + an exception list that hides the next real regression. Reading the one + function this PR owns is both honest and stricter: any name comparison + inside it fails, not just two spellings of one. + """ + import ast + + tree = ast.parse(ARGS_PATH.read_text(encoding="utf-8")) + fn = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "validate_algorithm_args" + ) + # `ast.unparse` drops comments and, once the docstring node is skipped, + # prose. Reading the raw text instead would flag the docstring's own + # description of the pattern it removed. + statements = fn.body[1:] if ast.get_docstring(fn) else fn.body + code = "\n".join(ast.unparse(node) for node in statements) + + assert "get_algorithm(" in code, "sanity: validation should still resolve the spec" + for line in code.splitlines(): + assert "advantage_estimator ==" not in line, f"name comparison in validate_algorithm_args: {line.strip()}" + assert "advantage_estimator in" not in line, f"name comparison in validate_algorithm_args: {line.strip()}" def test_validation_reads_spec_fields(): diff --git a/tests/algorithms/test_policy_loss_dispatch.py b/tests/algorithms/test_policy_loss_dispatch.py index c8fe6ecf2..d88798de5 100644 --- a/tests/algorithms/test_policy_loss_dispatch.py +++ b/tests/algorithms/test_policy_loss_dispatch.py @@ -164,9 +164,25 @@ def test_the_name_check_pattern_catches_every_spelling(): def test_both_paths_delegate_to_the_shared_estimator(): + """Delegation means the old kernels are gone, not that an import exists. + + Asserting only on the import line is satisfied by a file that imports the + shared handler and then ignores it -- exactly what a botched revert looks + like. The load-bearing half is the second assertion: if either path grows + its own call to a `ppo_utils` kernel again, the estimator is being chosen + somewhere other than the registry. + """ + direct_kernels = ( + "get_grpo_returns", + "get_reinforce_plus_plus_returns", + "get_reinforce_plus_plus_baseline_advantages", + "get_advantages_and_returns_batch", + ) 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" + assert "compute_advantages_and_returns" in src, f"{path.name} does not call the shared estimator" + for kernel in direct_kernels: + assert kernel not in src, f"{path.name} calls {kernel} directly instead of going through the registry" def test_loss_py_reads_kl_level_and_full_log_probs_from_the_spec(): From 74531487d086cd2a127c95c7030966aa6c2a3257 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 18:25:17 +0800 Subject: [PATCH 07/22] fix(algorithms): unbreak the CPU tests, and stop two more specs from lying Second review round. The first finding is a regression I introduced in this branch; the rest are gaps the previous round left. **`loss.py` became unimportable without a full Megatron install.** Importing `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY` from `relax.backends.megatron.data` pulls in `megatron.core.packed_seq_params` at module scope, and the CPU tests for `loss.py` stub `megatron.core` without that submodule. One import for one string cost two existing tests: `test_reinforce_plus_plus_wiring.py` and `test_reinforce_plus_plus_loss.py` went from 4 passed on main to 2 failed here. The full suite hid it -- run together, some earlier test installs a stub that makes the import succeed, so `pytest tests/` stayed green while `pytest tests/backends/megatron/test_reinforce_plus_plus_*.py` failed. CI that shards or selects tests would have caught what I did not. `loss.py` now spells the key out, with `test_loss_py_uses_the_canonical_mini_batch_key` comparing it against the definition as source text -- no import either way, and it works on a runner with no megatron at all. **PPO's spec claimed --fully-async support.** `validate_ppo_config` rejects both --fully-async and --hybrid for PPO, but the spec left the field at its default. This is the same class of error the previous commit fixed for the REINFORCE++ variants, in the one algorithm that commit did not touch. Declaring it also moves the failure from service registration to argument validation. **`advantage_normalization` accepted anything.** It is not a lookup key, so `_assert_spec_implementations_resolve` could not cover it -- and its failure mode is worse than a missing key: a typo does not raise, it selects the default branch in `loss.py` and trains with different maths. Now validated against its two legal values. **Step 1 arithmetic moved to float64**, matching `distributed_mean_std` and for the same reason. It also happens to fix an overflow the review found: components like [3e38, 2e38, 1e38, 0] overflowed while computing their own mean in float32 and tripped the combined-value guard even when their weight was zero. In float64 they standardise normally and the zero weight mutes them as it should. **Documented one thing float64 does not fix.** Two components summing to a constant should cancel exactly, and in float64 they do. But they arrive already cast to float32, and `C - r` does not survive that cast unchanged: the pair still sums to `C` (the addition rounds back) while the values have drifted, leaving ~1e-4 after standardisation, which step 3 divides by a batch std of the same order to produce advantages of order 1. A group with no signal gets a confident gradient whose direction is rounding noise. GRPO does not have this failure -- it standardises the sum, which is exactly constant. Recorded in examples/gdpo/README.md; a real fix needs a wider reward dtype end to end or a noise floor in step 3, neither of which belongs here. Tests: 1659 passed, same 2 failures + 2 errors as main@98a1274 here, and the two megatron CPU tests pass standalone again. --- examples/gdpo/README.md | 8 +++++ relax/algorithms/rewards.py | 31 +++++++++++++++++-- relax/algorithms/spec.py | 6 ++++ relax/backends/megatron/loss.py | 10 ++++-- .../filters/dynamic_sampling_filters.py | 8 +++++ relax/utils/arguments.py | 12 +++++++ tests/algorithms/test_algos_roles.py | 25 +++++++++++++++ .../algorithms/test_arguments_spec_driven.py | 20 +++++++++++- .../test_dispatch_parity_vs_main.py | 6 +++- 9 files changed, 119 insertions(+), 7 deletions(-) diff --git a/examples/gdpo/README.md b/examples/gdpo/README.md index 676bfa941..55261964f 100644 --- a/examples/gdpo/README.md +++ b/examples/gdpo/README.md @@ -81,9 +81,17 @@ GDPO_ARGS=( ## 已知偏差 1. **第三步的 batch 边界（已正确处理）**。调用方为效率会先合并多个训练批再调用 advantage，但 `_whiten_by_segment` 用 `mini_batch_sizes` 把它们切回**每个 optimizer 训练批各自白化**，因此 `num_rollout_minis > 1` 时仍对齐论文 Eq. 6，**不要求** `rollout_batch_size × n_samples_per_prompt == global_batch_size`。本脚本把 `4 × 8` 与 `--global-batch-size 32` 设成相等只是让例子最简单，并非必需。跨 DP 的 all-reduce 保证统计量覆盖全部 rank。**`--fully-async` 会在参数校验阶段被拒绝**——那条路径的切片可能小到只有一个样本，白化输出恒为 0。 + 2. **单个奖励时 GDPO 不等于 GRPO**。step1 除以 `std_g + 1e-4`、GRPO 除以 `std_g + 1e-6`，各组 `std_g` 不同 → 尺度因子逐组不同，step3 还会再做一次 batch 白化，所以不是「差一个正标量」那么简单。要 GRPO 语义就用 `--advantage-estimator grpo`。 + 3. **`--n-samples-per-prompt 2` 时幅度信息丢失**：任意两个不同值标准化后恒为 ±0.7071。示例用 8 就是为了避开这一点。 +4. **恒和分量在大量级下会产生数值假信号**。若两个分量恰好满足 `r₂ = C − r₁`，数学上应完全抵消为零（见上文「GDPO 不能做什么」）。但奖励在进入归一化前会被 cast 成 float32，而 `C − r` 这种值不一定能被 float32 精确表示——两者相加仍舍回 `C`（看起来恒和），各自却已偏离，标准化后留下约 `1e-4` 的残差。第三步再用同量级的 batch 标准差去除它，输出就变成 O(1) 的 advantage：**一个本无信号的组拿到了方向由舍入决定的梯度**。实测一组和为 `308.95172119140625` 的奖励，最终 advantage 为 `[-0.4344, 0.5251, -0.0908]`。 + + GRPO 没有这个问题：它归一化的是**和**，而和在 float32 下确实恒定，塌缩检测会直接置零。 + + 触发需要「恒和 + 分量绝对值远大于其组内差异」同时成立，实践中少见；彻底解决要么把奖励全链路加宽到 float64，要么在第三步引入噪声下限，两者都超出本 PR 范围。若你的奖励设计天然满足恒和，请直接用不等权重（此时不再抵消，也就不存在这个残差被放大的问题）。 + ## 冲突项 `--normalize-advantages`、`--custom-reward-post-process-path`、`--agentic-custom-advantage-path` 和 `--fully-async` 都不能与 GDPO 同用，参数校验阶段会直接报错。第一个会造成双重白化；第二、三个都会赶在归一化器之前从 `post_process_rewards` 返回，导致 GDPO 的前两步被静默跳过（这两个由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守）；第四个的统计窗口可能小到只剩一个样本。 diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index 59ce5f3d5..5736a60fb 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -228,12 +228,37 @@ def standardize_group_components(group: torch.Tensor) -> torch.Tensor: Each column is standardised on its own. A column that is flat across the group contributes exactly zero rather than ``0 / (0 + eps)`` noise, which is what lets the remaining columns keep their signal. + + The arithmetic runs in float64, for the same reason + :func:`relax.algorithms.numerics.distributed_mean_std` does: the means and + stds are the part where cancellation bites, and widening them is cheap. + + It does **not** fix the related failure worth knowing about. Two columns + that sum to a constant should standardise to exact opposites and cancel, + and in float64 they do -- exactly. But the columns arrive already cast to + float32 by :func:`extract_reward_components`, and a value like + ``C - r`` does not survive that cast unchanged. The pair still *sums* to + ``C`` in float32 (the addition rounds back), so the group looks flat, while + the individual values have drifted enough to leave a residue of order 1e-4 + after standardisation. Step 3 then divides that residue by a batch std of + the same order and returns advantages of order 1: a group with no signal + gets a confident gradient whose direction is decided by rounding. + + Measured on rewards summing to 308.95172119140625, the batch comes out as + [-0.4344, 0.5251, -0.0908]. GRPO does not have this failure, because it + standardises the sum, which *is* exactly constant, and its collapse check + catches it. Raising the precision here does not help -- the information was + lost before this function saw it. Documented in examples/gdpo/README.md + under the deviations; fixing it needs either a wider reward dtype end to + end or a noise floor in step 3, neither of which belongs in this PR. """ - centered = group - group.mean(dim=0, keepdim=True) - std = group.std(dim=0) + work = group.double() + centered = work - work.mean(dim=0, keepdim=True) + std = work.std(dim=0) collapsed = collapsed_columns(group, dim=0) scaled = centered / (std + GDPO_EPS) - return torch.where(collapsed.unsqueeze(0), torch.zeros_like(scaled), scaled) + scaled = torch.where(collapsed.unsqueeze(0), torch.zeros_like(scaled), scaled) + return scaled.to(group.dtype) def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: diff --git a/relax/algorithms/spec.py b/relax/algorithms/spec.py index 446a214a7..7a4c79bf6 100644 --- a/relax/algorithms/spec.py +++ b/relax/algorithms/spec.py @@ -251,6 +251,12 @@ def is_group_normalized(self) -> bool: advantage_fn="gae", policy_loss_fn="ppo_clip", needs_critic=True, + # `validate_ppo_config` in `relax/utils/training/ppo_utils.py` rejects + # both --fully-async and --hybrid for PPO. Declaring it means argument + # validation says so up front instead of the Controller saying it at + # service-registration time, and stops the spec from asserting the + # opposite by omission. + supports_fully_async=False, ), "reinforce_plus_plus": AlgorithmSpec( name="reinforce_plus_plus", diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index fc7aafe64..f96462cb6 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -12,7 +12,6 @@ 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.backends.megatron.data import ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY 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 ( @@ -601,7 +600,14 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) # rollout. Passing them is what lets a batch-level statistic describe one # batch instead of the whole merged rollout; estimators that do not have # one absorb this in **_unused. - mini_batch_sizes=rollout_data.get(ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY), + # The key is spelled out rather than imported from + # `relax.backends.megatron.data`: that module pulls in + # `megatron.core.packed_seq_params` at import time, and the CPU-only + # tests for this file stub `megatron.core` without that submodule, so + # importing it for one string makes `loss.py` unimportable for them. + # `test_loss_py_uses_the_canonical_mini_batch_key` pins this literal to + # the definition so the two cannot drift. + mini_batch_sizes=rollout_data.get("rollout_mini_local_sample_counts"), ) # Optional pure OPD mode: remove all non-OPD reward contribution. diff --git a/relax/engine/filters/dynamic_sampling_filters.py b/relax/engine/filters/dynamic_sampling_filters.py index b2425b359..5fe387b90 100644 --- a/relax/engine/filters/dynamic_sampling_filters.py +++ b/relax/engine/filters/dynamic_sampling_filters.py @@ -18,6 +18,14 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): multi-reward algorithm it is whether *every* component is flat; judging those by the summed scalar would drop the groups the algorithm exists to keep. + + Note this can now raise rather than merely returning a verdict: for a + multi-reward algorithm it runs the same component extraction the reward + stage does, which rejects malformed rewards. `call_dynamic_filter` does not + catch, so a broken reward function fails the rollout here instead of a few + lines later in `post_process_rewards`. That is the same failure, earlier and + attributed to the group that caused it -- but it does mean this filter is no + longer guaranteed to be side-effect-free on bad data. """ keep = group_carries_reward_signal(args, samples) return DynamicFilterOutput( diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 95117ee77..a0f8c11e6 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -2999,6 +2999,18 @@ def _assert_spec_implementations_resolve(spec) -> None: f"Available: {sorted(table)}." ) + # Not a lookup key, so the loop above cannot cover it, but it has the same + # failure mode and a worse consequence: a typo does not raise, it silently + # selects the default branch in `relax/backends/megatron/loss.py` and + # changes the training maths. + valid_normalizations = {"whiten", "token_global"} + if spec.advantage_normalization not in valid_normalizations: + raise ValueError( + f"Algorithm {spec.name!r} declares advantage_normalization=" + f"{spec.advantage_normalization!r}, which is not a known mode. " + f"Available: {sorted(valid_normalizations)}." + ) + def validate_algorithm_args(args) -> None: """Apply the constraints the algorithm registry declares for this run. diff --git a/tests/algorithms/test_algos_roles.py b/tests/algorithms/test_algos_roles.py index e21f2ef3c..f553eb335 100644 --- a/tests/algorithms/test_algos_roles.py +++ b/tests/algorithms/test_algos_roles.py @@ -133,3 +133,28 @@ 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" + + +def test_loss_py_uses_the_canonical_mini_batch_key(): + """`loss.py` spells the key out; this pins it to the definition. + + Importing `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY` from + `relax.backends.megatron.data` would be the obvious way to avoid a + duplicated literal, but that module imports + `megatron.core.packed_seq_params` at module scope, and the CPU-only tests + for `loss.py` stub `megatron.core` without it -- one import for one string + made `loss.py` unimportable for them. Comparing the source text keeps the + two in sync without either file importing the other, and works on a runner + with no megatron at all. + """ + import re + + root = pathlib.Path(__file__).resolve().parents[2] + data_src = (root / "relax" / "backends" / "megatron" / "data.py").read_text(encoding="utf-8") + match = re.search(r'ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY\s*=\s*"([^"]+)"', data_src) + assert match, "the canonical key definition moved or changed shape" + + loss_src = (root / "relax" / "backends" / "megatron" / "loss.py").read_text(encoding="utf-8") + assert f'"{match.group(1)}"' in loss_src, ( + f"loss.py no longer reads the canonical key {match.group(1)!r}; the literal drifted" + ) diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py index 2df3aded1..eddbaefa6 100644 --- a/tests/algorithms/test_arguments_spec_driven.py +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -472,7 +472,7 @@ def test_supports_fully_async_is_declared_wherever_it_is_false(): from relax.algorithms import get_algorithm, list_algorithm_names unsupported = {n for n in list_algorithm_names() if not get_algorithm(n).supports_fully_async} - assert unsupported == {"gdpo", "reinforce_plus_plus", "reinforce_plus_plus_baseline"} + assert unsupported == {"gdpo", "ppo", "reinforce_plus_plus", "reinforce_plus_plus_baseline"} def test_dynamic_sampling_filter_warns_for_multi_reward(arguments_module, caplog): @@ -525,3 +525,21 @@ def test_yaml_cannot_bypass_the_frozen_reinforce_plus_plus_contract(arguments_mo ) with pytest.raises(ValueError, match="custom-reward-post-process-path"): arguments_module.apply_custom_config_overrides(args) + + +def test_a_typo_in_advantage_normalization_is_rejected(arguments_module): + """It is not a lookup key, so a typo would otherwise pass validation. + + `advantage_normalization` selects between masked whitening and + REINFORCE++'s token-global normalisation in `loss.py`. An unrecognised + value there does not raise -- it falls through to the default branch and + quietly trains with different maths, which is the failure this whole + registry exists to remove. + """ + import dataclasses + + from relax.algorithms import get_algorithm + + broken = dataclasses.replace(get_algorithm("grpo"), advantage_normalization="typo") + with pytest.raises(ValueError, match="advantage_normalization"): + arguments_module._assert_spec_implementations_resolve(broken) diff --git a/tests/algorithms/test_dispatch_parity_vs_main.py b/tests/algorithms/test_dispatch_parity_vs_main.py index d271af6d4..d8deb27d5 100644 --- a/tests/algorithms/test_dispatch_parity_vs_main.py +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -548,7 +548,11 @@ def test_loss_py_actually_forwards_the_mini_batch_boundaries(): call = re.search(r"compute_advantages_and_returns_impl\((.*?)\n \)", src, re.DOTALL) assert call, "compute_advantages_and_returns_impl call not found" - assert "mini_batch_sizes=rollout_data.get(ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY)" in call.group(1), ( + # The key is a literal in loss.py, not the imported constant: importing it + # from `megatron.data` drags in `megatron.core.packed_seq_params` and makes + # loss.py unimportable for its own CPU tests. `test_algos_roles.py` pins the + # literal to the definition. + assert 'mini_batch_sizes=rollout_data.get("rollout_mini_local_sample_counts")' in call.group(1), ( "loss.py must pass the per-training-batch counts; without them GDPO's step 3 " "normalises over the whole merged rollout again" ) From 6f214479844f1505c1479431481d574c7354a5bd Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Wed, 12 Aug 2026 18:34:54 +0800 Subject: [PATCH 08/22] docs(gdpo): state the motivation as one formula instead of two half-truths The second review round showed the corrected motivation was still wrong, in the opposite direction from the original. Both errors turn out to be the same formula read at two extremes, so this replaces them with the formula. The combined advantage is `sum_k w_k * z_k` with each `z_k` at unit variance, so its variance is `sum_k w_k^2 + 2 sum_{i 2.00x, 0 -> 1.41x, -0.5 -> 1.00x, -1 -> 0.00x. The scale-disparity half of the motivation was checked again and stands unchanged: GRPO's direction is dominated by whichever component has the largest spread, because it standardises the raw sum. No code change -- `normalize_gdpo_decoupled` already computed this correctly. What was wrong was every sentence describing it. --- docs/en/examples/algorithms.md | 4 +++- docs/zh/examples/algorithms.md | 4 +++- examples/gdpo/README.md | 23 ++++++++++++++++------- relax/algorithms/rewards.py | 34 +++++++++++++++++++++------------- 4 files changed, 43 insertions(+), 22 deletions(-) diff --git a/docs/en/examples/algorithms.md b/docs/en/examples/algorithms.md index 3f6623523..1e70e541f 100644 --- a/docs/en/examples/algorithms.md +++ b/docs/en/examples/algorithms.md @@ -301,7 +301,9 @@ $$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\ma **Why this beats GRPO:** summing first and standardizing once discards two things. -*Relative strength between groups.* Per-group standardization forces every group to unit variance, so a group where only one component varies and a group where both vary come out identical. Standardizing each component first makes the latter twice the amplitude, and step 3 whitens across the *batch*, so the difference survives into the final advantage. Measured (G=2, two groups in one batch): GRPO gives both groups ±0.707; GDPO gives ±0.548 and ±1.095. +*Correlation structure between components.* GRPO standardizes the sum, so every group comes out at unit variance whether its components corroborate or contradict each other. GDPO gives each component unit variance first, so the combined variance is `Σwᵢ² + 2Σwᵢwⱼρᵢⱼ` — `2 + 2ρ` for two equal weights — and is therefore **decided by the correlation**: `ρ→+1` amplifies (measured 2x), `ρ=0` gives √2, `ρ→−1` attenuates to zero. Step 3 whitens across the *batch*, so that between-group difference reaches the final advantage. Measured (G=2, ρ=+1, two groups in one batch): GRPO gives both ±0.707; GDPO gives ±0.548 and ±1.095. + +This is **not** "more varying components means more signal" — that holds only for `ρ>0`. At `ρ=−0.8` the combined signal is 0.63x a single component. *Scale disparity between components.* A `correctness` in {0, 1} added to a reward in the hundreds (the paper's maths setup scores response length) yields a sum whose variance is essentially the large component's, so GRPO's direction is decided by it alone. GDPO gives each component unit variance first, so a weight expresses relative importance rather than units. diff --git a/docs/zh/examples/algorithms.md b/docs/zh/examples/algorithms.md index 5c5d26b57..612e4ac00 100644 --- a/docs/zh/examples/algorithms.md +++ b/docs/zh/examples/algorithms.md @@ -298,7 +298,9 @@ $$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\ma **相对 GRPO 的收益**：GRPO 把各分量相加后只做一次组内标准化，这会丢掉两类信息。 -*一是组间的相对强度*：组内标准化把每一组都拉到单位方差，于是「只有一个分量在变」的组与「两个分量都在变」的组得到完全相同的 advantage。GDPO 让各分量先各自标准化，后者幅度自然是两倍；而第三步的 batch 白化是**跨组**的，这个差异会保留到最终 advantage。实测（G=2，一个 batch 两组）：GRPO 两组同为 ±0.707，GDPO 分别为 ±0.548 与 ±1.095。 +*一是分量间的相关结构*：GRPO 对和做组内标准化，每组输出恒为单位方差，无论分量彼此印证还是互相矛盾。GDPO 让各分量先各自标准化再加权求和，组合方差为 `Σwᵢ² + 2Σwᵢwⱼρᵢⱼ`（等权重两分量时为 `2 + 2ρ`），**由相关系数决定**：`ρ→+1` 增强（实测 2 倍）、`ρ=0` 为 √2 倍、`ρ→−1` 减弱直至归零。第三步的 batch 白化跨组进行，这个组间强度差异会保留到最终 advantage。实测（G=2、ρ=+1，一个 batch 两组）：GRPO 两组同为 ±0.707，GDPO 分别为 ±0.548 与 ±1.095。 + +注意这条**不是**「变化的分量越多信号越强」——那只在 `ρ>0` 时成立。`ρ=−0.8` 时组合信号反而只有单分量的 0.63 倍。 *二是分量间的尺度差异*：`correctness ∈ {0,1}` 与一个取值上百的分量（论文实验用响应长度）相加时，和的方差几乎全部来自后者，GRPO 的方向由它单独决定；GDPO 先让每个分量单位方差，权重才真正表达相对重要性而非量纲。 diff --git a/examples/gdpo/README.md b/examples/gdpo/README.md index 55261964f..d627d9575 100644 --- a/examples/gdpo/README.md +++ b/examples/gdpo/README.md @@ -13,20 +13,29 @@ GRPO 把它们**加起来**再做一次组内归一化。这一步会丢掉两类信息。 -**一、组间的相对强度。** 组内标准化把每一组的 advantage 都拉到单位方差，于是「只有 correctness 在变」的组和「两个分量都在变」的组，得到完全相同的 advantage。GDPO 对每个分量各自标准化后相加，两个分量都起作用时幅度自然是两倍；而第三步的 batch 白化是**跨组**的，所以这个差异会一路保留到最终 advantage： +**一、分量之间的相关结构。** GRPO 对和做组内标准化，每一组的输出恒为单位方差——无论这组的两个分量是彼此印证还是互相矛盾。GDPO 先让每个分量各自单位方差再加权求和，组合结果的方差是 -| | 组 A（只有 correctness 变化） | 组 B（两个分量都变化） | -| ---------------- | ----------------------------- | ---------------------- | -| GRPO | ±0.707 | ±0.707（分不出） | -| GDPO（含第三步） | ±0.548 | ±1.095（B 强一倍） | +``` +Var = Σ wᵢ² + 2 Σᵢ<ⱼ wᵢwⱼ ρᵢⱼ 等权重两分量时 = 2 + 2ρ +``` + +**由分量间的相关系数 ρ 决定**： + +| ρ | 含义 | 组合信号 | +| ---- | ---------------------- | ----------------- | +| → +1 | 两个分量指向同一批样本 | 增强（实测 2 倍） | +| 0 | 互不相关 | √2 倍 | +| → −1 | 两个分量互相矛盾 | 减弱，极限时归零 | + +第三步的 batch 白化是**跨组**的，所以这个组间强度差异会保留到最终 advantage。GRPO 看到的则是所有组都一样强。 **二、分量之间的尺度差异。** `correctness ∈ {0,1}` 与一个取值上百的 `format`（或论文实验里的响应长度）相加时，和的方差几乎全部来自大尺度那一维，GRPO 的方向就由它单独决定。GDPO 先让每个分量单位方差，权重才真正表达「相对重要性」而不是量纲。极端一点：`correctness=[1,1,0,0]`、`format=[0,100,200,300]`（两者排序相反）时，GRPO 给答错的长响应最高 advantage，GDPO 不会。 ## GDPO 不能做什么 -**各分量在组内恰好加和为常数时，GDPO 也救不回来。** 若 `correctness + format ≡ C`，则 `format = C − correctness`，两者标准化后恒有 `z_format = −z_correctness`，**等权重下完全抵消为零**——与 GRPO 得到同样的结果。 +**各分量在组内恰好加和为常数时，GDPO 也救不回来。** 这正是上表 `ρ = −1` 那一行：若 `correctness + format ≡ C`，则 `format = C − correctness`，两者标准化后恒有 `z_format = −z_correctness`，**等权重下完全抵消为零**——与 GRPO 得到同样的结果。 -这一点本文档此前写反了，说这正是 GDPO 的优势场景。它不是：那是一个数学上不可能被等权重 GDPO 区分的情形。只有不等权重（如 `--gdpo-reward-weights 2.0 1.0`）能在这类组上拿到信号。 +这一点本文档此前写反了，说这正是 GDPO 的优势场景。它不是：那是 `ρ = −1`，是组合方差 `2 + 2ρ` 恰好取零的极端。只有不等权重（如 `--gdpo-reward-weights 2.0 1.0`）能打破这个平衡。 （另外，若两个分量在组内**都**恒定，GDPO 与 GRPO 一样返回零，不会无中生有。） diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index 5736a60fb..04da50be2 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -271,19 +271,27 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl :func:`relax.algorithms.advantages.advantage_gdpo`. What standardising per component actually buys, stated carefully because - it is easy to overclaim: the combined advantage is ``sum_k w_k * z_k``, - where each ``z_k`` has unit variance within the group. GRPO instead - standardises ``sum_k r_k``, so a component with a large spread dominates - the direction. When the components have comparable spread the two agree up - to a positive scalar; they diverge when the spreads differ, which is the - case GDPO is for -- a correctness reward in ``{0, 1}`` combined with a - length reward in the hundreds is decided almost entirely by length under - GRPO, and half by each under GDPO. - - It does **not** rescue a group whose components sum to a constant. There - ``r_2 = C - r_1`` forces ``z_2 = -z_1``, so equal weights cancel to exactly - zero -- the same answer GRPO gives. Unequal weights break the tie; equal - weights cannot. + it is easy to overclaim in two opposite directions (this docstring has + managed both): + + The combined advantage is ``sum_k w_k * z_k`` with each ``z_k`` at unit + variance, so its variance is ``sum_k w_k^2 + 2 sum_{i +1``) amplify, contradicting ones (``rho -> -1``) attenuate and + at the limit cancel. + + Both claims this file got wrong earlier are special cases of that one + formula. 'Equal sums are rescued by GDPO' is ``rho = -1`` -- they cancel, + GDPO included. 'Two varying components mean a stronger signal' is + ``rho = +1`` -- true there, false at ``rho = -0.8``, where the combination + measures 0.63x a single component. + + Separately and unconditionally: GRPO's direction is dominated by whichever + component has the largest spread, because it standardises the raw sum. A + correctness reward in ``{0, 1}`` added to a length reward in the hundreds + is decided almost entirely by length under GRPO, and half by each here. """ keys = resolve_gdpo_keys(args) weights = resolve_gdpo_weights(args, keys) From c1399bf3e285a43f69e6bdd7df38d4f04e841797 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Thu, 13 Aug 2026 13:19:16 +0800 Subject: [PATCH 09/22] fix(gdpo): stop cancelling components from becoming a gradient MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A third adversarial review round found two defects that two prior rounds missed, both outside the GDPO maths those rounds concentrated on. Two components summing to a constant should standardise to exact opposites and cancel. Instead the ill-conditioned division left a residue that step 3 divided by the std of that very residue. Measured on rewards summing to 308.95172119140625, a group carrying no signal came out at `[-0.5770, 1.1539, -0.5770]` — finite, plausible, and pointing wherever the last bits of the reward happened to fall. Two changes, and it is worth separating them because the tempting summary ("the floor stops a fake gradient") is not what happens: - `extract_reward_components` builds float64. This is the half that matters: the residue drops from 1e-1 (larger than the signal) to 1e-10, and `GDPO_EPS` in step 3's denominator caps the amplification at ~1e-6. - `combine_group` applies a noise floor, so cancellation is *exactly* zero. 1e-6 is not zero, and consumers test for zero: the dynamic-sampling filter would keep a group contributing nothing. The threshold is computed, not tuned — `component_noise_scale` returns `eps * max|x| / (std + GDPO_EPS)`, the condition number of the standardisation, which is fifteen orders below the signal on well-conditioned rewards and matches the observed residue (bound 5.1e-10, measured 4.1e-10) on the pathological one. This was previously documented as an unfixable deviation. That was wrong on the diagnosis as well as the priority: it is not a precision limit, it is a gap in the collapse check. `relax/distributed/ray/rollout.py` did not filter `reward is None` while its twin in `relax/agentic/rollout.py` did, so the shared helper built a float32 tensor from a None and raised `TypeError` — from a logging helper, on the rollout's way out. `reward is None` is reachable: under `--group-rm` the group reward is assigned in one shot that is skipped entirely when the rollout aborts, which is why `sglang_rollout.py`'s "reward is not None" assert exempts `group_rm` in the first place. All single-reward algorithms were affected, including the default. Separately, eval may run a different reward model (`EvalConfig.rm_type`), so an eval reward legitimately need not carry `--gdpo-reward-keys`. Metrics were enforcing a training contract that is not in force there, and failing eval for it. Both decisions now live in `metrics_group_verdict`, which returns a tri-state and is importable without `sglang` — the two copies had drifted apart once and neither was testable, let alone tested. `_whiten_by_segment` treated `mini_batch_sizes=None` as "use one window". That is not a coarser version of the same objective: the repository's own `test_merging_the_batches_would_flip_signs_not_just_rescale` shows half the advantages changing sign, with every metric finite. It now raises, through the existing MAX all-reduce so the whole group fails together. `--gdpo-reward-weights` is checked for non-finite and all-zero-in-float32 values during argument validation rather than at the first rollout. - `advantage_normalization` had no test pinning its value — dropping `"token_global"` from a REINFORCE++ spec, or flipping the comparison in `loss.py`, left the suite green. Both call sites and the exact set are now transcribed from main. - `advantage_gae` had only a `co_names` check. It is now compared numerically against a transcription of main's inline branch. PPO is the algorithm this refactor changed most and the one with no GPU smoke. - `test_grouping_uses_group_index` only ever built a contiguous layout, so a position-based implementation satisfied it. It now interleaves two groups. - Every guard above was mutation-tested: each one, removed, turns a test red. The previous round's corrections landed in the main implementation and left the periphery behind. The smoke script header still carried the motivation retracted in 9f1e02e; the filter docstring described "every component is flat" where the code asks whether the combination cancels; a test docstring cited a `--global-batch-size` the shipped example does not use; the overflow comment described a mean that no longer overflows; `.detach()` was called "load-bearing" when both call sites pass lists and never reach it. `tests/algorithms` 745 passed. Full suite 1686 passed, with the same 2 failures and 2 errors present on the base 98a1274 (verified in a detached worktree). `pre-commit run --all-files` clean. --- docs/en/examples/algorithms.md | 2 +- docs/zh/examples/algorithms.md | 2 +- examples/gdpo/README.md | 15 +- examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh | 15 +- relax/agentic/rollout.py | 15 +- relax/algorithms/advantages.py | 39 ++- relax/algorithms/rewards.py | 253 ++++++++++++++---- relax/backends/megatron/loss.py | 9 +- relax/distributed/ray/rollout.py | 15 +- .../filters/dynamic_sampling_filters.py | 14 +- relax/utils/arguments.py | 21 ++ tests/algorithms/test_advantage_estimators.py | 5 + .../algorithms/test_arguments_spec_driven.py | 38 ++- .../test_dispatch_parity_vs_main.py | 136 ++++++++++ .../algorithms/test_distributed_whitening.py | 15 +- tests/algorithms/test_gdpo.py | 179 ++++++++++++- .../algorithms/test_multi_reward_consumers.py | 91 ++++++- 17 files changed, 756 insertions(+), 108 deletions(-) diff --git a/docs/en/examples/algorithms.md b/docs/en/examples/algorithms.md index 1e70e541f..93ccaf48d 100644 --- a/docs/en/examples/algorithms.md +++ b/docs/en/examples/algorithms.md @@ -342,7 +342,7 @@ A complete runnable example lives in [`examples/gdpo/`](https://github.com/redai Two differences between this implementation and the paper. Confirm they are acceptable before training. Step 3's batch boundary used to be a third; it has since been corrected — see below. -**Step 3's batch boundary (now aligned).** Eq. 6 normalises over one training batch. The caller merges `num_rollout_minis` of them with `concat_rollout_batches` before the advantage stage, so step 3 has to be told where the boundaries are. They travel in `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY`, which all three actor paths (colocate and hybrid) set; `loss.py` forwards them to the advantage dispatcher as `mini_batch_sizes`, and **only GDPO reads it** — every other estimator absorbs it in `**_unused` and is bit-identical either way. Each segment all-reduces across the data-parallel group, so the statistics cover both a whole training batch and every rank. Why it matters: whitening merged batches centres them all on a pooled mean, and on a measured example four of eight samples **change sign** — a different objective, not a precision difference. +**Step 3's batch boundary (now aligned).** Eq. 6 normalises over one training batch. The caller merges `num_rollout_minis` of them with `concat_rollout_batches` before the advantage stage, so step 3 has to be told where the boundaries are. They travel in `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY`, which all three actor paths (colocate and hybrid) set; `loss.py` forwards them to the advantage dispatcher as `mini_batch_sizes`, and **only GDPO reads it** — every other estimator absorbs it in `**_unused` and is bit-identical either way. Each segment all-reduces across the data-parallel group, so the statistics cover both a whole training batch and every rank. Why it matters: whitening merged batches centres them all on a pooled mean, and on a measured example four of eight samples **change sign** — a different objective, not a precision difference. That is also why absent boundaries are an **error** rather than a fallback to merged whitening: a caller that omits the metadata would optimise the wrong objective with loss and grad_norm both fine. **`--fully-async` remains unsupported** and is rejected during argument validation: it hands advantage computation to the single-replica Advantages deployment, which has no data-parallel group, never sees the batch boundaries, and consumes one `global_batch_size / num_iters_per_train_update` slice at a time — when that quotient is 1 the whitened output is identically zero and the run trains on no signal at all, quietly. diff --git a/docs/zh/examples/algorithms.md b/docs/zh/examples/algorithms.md index 612e4ac00..f1e38237b 100644 --- a/docs/zh/examples/algorithms.md +++ b/docs/zh/examples/algorithms.md @@ -339,7 +339,7 @@ GDPO_ARGS=( 以下两点是实现与论文之间的实际差异，训练前请确认可以接受。第三步的 batch 边界曾经也在此列，现已修正——见下。 -**第三步的 batch 边界（已对齐）**。论文 Eq. 6 在**一个训练批**上归一化。调用方会先把 `num_rollout_minis` 个训练批用 `concat_rollout_batches` 合并再进 advantage 阶段，所以第三步必须知道批边界。边界由 `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY` 携带（colocate 与 hybrid 三条路径都写入），`loss.py` 作为 `mini_batch_sizes` 传给 advantage 分发器，**只有 GDPO 消费**——其余估计器由 `**_unused` 吞掉，逐位不变。每段各自跨 DP all-reduce，所以统计量既覆盖完整训练批、也覆盖全部 rank。这一点为什么重要：合并白化会把两个批都对着共同均值中心化，实测 8 个样本里有 4 个**符号翻转**——那是另一个优化目标，不是精度差异。 +**第三步的 batch 边界（已对齐）**。论文 Eq. 6 在**一个训练批**上归一化。调用方会先把 `num_rollout_minis` 个训练批用 `concat_rollout_batches` 合并再进 advantage 阶段，所以第三步必须知道批边界。边界由 `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY` 携带（colocate 与 hybrid 三条路径都写入），`loss.py` 作为 `mini_batch_sizes` 传给 advantage 分发器，**只有 GDPO 消费**——其余估计器由 `**_unused` 吞掉，逐位不变。每段各自跨 DP all-reduce，所以统计量既覆盖完整训练批、也覆盖全部 rank。这一点为什么重要：合并白化会把两个批都对着共同均值中心化，实测 8 个样本里有 4 个**符号翻转**——那是另一个优化目标，不是精度差异。正因如此，边界缺失时 GDPO **直接报错**而不是退回合并白化：一个漏写这份元数据的调用方会在 loss、grad_norm 全部正常的情况下优化错误的目标。 **`--fully-async` 仍不受支持**，参数校验阶段直接拒绝：那条路径把 advantage 计算交给单副本的 Advantages 服务，它没有数据并行通信域，也拿不到批边界，且每次只消费 `global_batch_size / num_iters_per_train_update` 的一个切片；当这个商为 1 时白化输出恒为 0，训练会安静地在零信号上跑完。 diff --git a/examples/gdpo/README.md b/examples/gdpo/README.md index d627d9575..ec76629f9 100644 --- a/examples/gdpo/README.md +++ b/examples/gdpo/README.md @@ -89,17 +89,24 @@ GDPO_ARGS=( ## 已知偏差 -1. **第三步的 batch 边界（已正确处理）**。调用方为效率会先合并多个训练批再调用 advantage，但 `_whiten_by_segment` 用 `mini_batch_sizes` 把它们切回**每个 optimizer 训练批各自白化**，因此 `num_rollout_minis > 1` 时仍对齐论文 Eq. 6，**不要求** `rollout_batch_size × n_samples_per_prompt == global_batch_size`。本脚本把 `4 × 8` 与 `--global-batch-size 32` 设成相等只是让例子最简单，并非必需。跨 DP 的 all-reduce 保证统计量覆盖全部 rank。**`--fully-async` 会在参数校验阶段被拒绝**——那条路径的切片可能小到只有一个样本，白化输出恒为 0。 +1. **第三步的 batch 边界（已正确处理）**。调用方为效率会先合并多个训练批再调用 advantage，但 `_whiten_by_segment` 用 `mini_batch_sizes` 把它们切回**每个 optimizer 训练批各自白化**，因此 `num_rollout_minis > 1` 时仍对齐论文 Eq. 6，**不要求** `rollout_batch_size × n_samples_per_prompt == global_batch_size`。本脚本把 `4 × 8` 与 `--global-batch-size 32` 设成相等只是让例子最简单，并非必需——把 `--global-batch-size` 改成 16 就会跑出两段。跨 DP 的 all-reduce 保证统计量覆盖全部 rank。**`--fully-async` 会在参数校验阶段被拒绝**——那条路径的切片可能小到只有一个样本，白化输出恒为 0。 + + 调用方**必须**提供这份切分（`rollout_mini_local_sample_counts`）。缺失时 GDPO 直接报错，不回退到「整个 rollout 白化一次」：那不是同一个目标的粗糙版本，而是另一个目标——`test_merging_the_batches_would_flip_signs_not_just_rescale` 里 8 个样本有 4 个符号翻转，而 loss、grad_norm、advantage 均值全都正常。 2. **单个奖励时 GDPO 不等于 GRPO**。step1 除以 `std_g + 1e-4`、GRPO 除以 `std_g + 1e-6`，各组 `std_g` 不同 → 尺度因子逐组不同，step3 还会再做一次 batch 白化，所以不是「差一个正标量」那么简单。要 GRPO 语义就用 `--advantage-estimator grpo`。 3. **`--n-samples-per-prompt 2` 时幅度信息丢失**：任意两个不同值标准化后恒为 ±0.7071。示例用 8 就是为了避开这一点。 -4. **恒和分量在大量级下会产生数值假信号**。若两个分量恰好满足 `r₂ = C − r₁`，数学上应完全抵消为零（见上文「GDPO 不能做什么」）。但奖励在进入归一化前会被 cast 成 float32，而 `C − r` 这种值不一定能被 float32 精确表示——两者相加仍舍回 `C`（看起来恒和），各自却已偏离，标准化后留下约 `1e-4` 的残差。第三步再用同量级的 batch 标准差去除它，输出就变成 O(1) 的 advantage：**一个本无信号的组拿到了方向由舍入决定的梯度**。实测一组和为 `308.95172119140625` 的奖励，最终 advantage 为 `[-0.4344, 0.5251, -0.0908]`。 +4. **恒和分量的数值假信号（已修复，记录在此因为它解释了两处实现选择）**。若两个分量恰好满足 `r₂ = C − r₁`，数学上应完全抵消为零（见上文「GDPO 不能做什么」）。这里的标准化在这种输入下是**病态**的：它要除以一个接近零的 std，相对误差被放大 `max|x| / std` 倍。 - GRPO 没有这个问题：它归一化的是**和**，而和在 float32 下确实恒定，塌缩检测会直接置零。 + 奖励曾经在进入归一化前被 cast 成 float32，而 `C − r` 这种值不一定能被 float32 精确表示——两者相加仍舍回 `C`（看起来恒和），各自却已偏离。残差量级约 `1e-1`，**比信号本身还大**，第三步再除以同量级的 batch 标准差，输出就是 O(1) 的 advantage：一个本无信号的组拿到了方向由舍入决定的梯度。实测一组和为 `308.95172119140625` 的奖励，最终 advantage 为 `[-0.5770, 1.1539, -0.5770]`。 + + 现在两处一起挡住它，各管一半： - 触发需要「恒和 + 分量绝对值远大于其组内差异」同时成立，实践中少见；彻底解决要么把奖励全链路加宽到 float64，要么在第三步引入噪声下限，两者都超出本 PR 范围。若你的奖励设计天然满足恒和，请直接用不等权重（此时不再抵消，也就不存在这个残差被放大的问题）。 + - **分量全链路用 float64**（`extract_reward_components`）。残差降到 `1e-10` 量级，第三步分母里的 `GDPO_EPS = 1e-4` 又钳住了放大倍数，最坏输出已经只有 `1e-6`。**假梯度问题到这一步就没了。** + - **`combine_group` 的噪声地板**。`1e-6` 毕竟不是零，而下游好几处是按「等于零」判断的——filter 会保留一个不贡献梯度的组，「整批无信号」的告警也永远不会触发。地板把它归成精确的零。阈值由 `component_noise_scale` 从 `eps · max|x| / (std + GDPO_EPS)` 算出，不是拍的：良态输入下它比信号低十五个数量级。 + + GRPO 没有这个问题：它归一化的是**和**，而和在 float32 下确实恒定，塌缩检测会直接置零。 ## 冲突项 diff --git a/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh index 5c9a04442..66856dc2d 100644 --- a/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh +++ b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh @@ -5,12 +5,15 @@ # Qwen3-0.6B single-GPU GDPO training on GSM8K. # # GDPO (arXiv 2601.05242) standardizes each reward component within its prompt -# group before combining them. That keeps two things GRPO's single -# standardization of the summed reward loses: the relative strength between -# groups (per-group standardization forces every group to unit variance), and -# the balance between components of very different scale. It does NOT rescue a -# group whose components sum to a constant -- those cancel to zero under equal -# weights, same as GRPO. See examples/gdpo/README.md. +# group before combining them. The combined advantage is sum_k w_k * z_k with +# each z_k at unit variance, so its variance is 2 + 2*rho for two equal +# weights: what GDPO preserves is the CORRELATION between components, which +# GRPO destroys by standardizing the summed reward to unit variance for every +# group regardless. Corroborating components amplify, contradicting ones +# attenuate, and at rho = -1 they cancel -- so it does NOT rescue a group whose +# components sum to a constant. Separately and unconditionally, GRPO's +# direction is dominated by whichever component has the largest raw spread; +# here each contributes according to its weight. See examples/gdpo/README.md. # The reward function in reward_gdpo.py returns both components; # --gdpo-reward-keys names them. # diff --git a/relax/agentic/rollout.py b/relax/agentic/rollout.py index ec37e392a..480ee0d7e 100644 --- a/relax/agentic/rollout.py +++ b/relax/agentic/rollout.py @@ -20,7 +20,7 @@ ) from relax.agentic.profile import TRACE_KEY from relax.algorithms import get_algorithm -from relax.algorithms.rewards import group_carries_reward_signal +from relax.algorithms.rewards import metrics_group_verdict from relax.engine.filters.base_types import MetricGatherer, call_dynamic_filter from relax.engine.rollout import on_policy_distillation as opd from relax.engine.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput @@ -1301,12 +1301,17 @@ def _compute_zero_std_metrics(args, all_samples: list[Sample]) -> dict[str, floa all_sample_groups = group_by(all_samples, lambda sample: sample.group_index) interesting_rewards = [] for group in all_sample_groups.values(): - rewarded = [sample for sample in group if sample.reward is not None] # Counted as flat only when it carries no signal *for this algorithm*: - # for a multi-reward one that means every component is flat, not that - # the --reward-key scalar happens to be. - if not rewarded or group_carries_reward_signal(args, rewarded): + # for a multi-reward one that means the weighted combination cancels, + # not that the --reward-key scalar happens to be flat. + # + # `is not True`, so a group that is unscored or unreadable is skipped + # rather than counted -- which is what this metric did before, since it + # dropped `reward is None` samples and then required a non-empty list. + # The distributed copy counts them instead; see `metrics_group_verdict`. + if metrics_group_verdict(args, group) is not True: continue + rewarded = [sample for sample in group if sample.reward is not None] interesting_rewards.append(str(round(rewarded[0].get_reward_value(args), 1))) return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} diff --git a/relax/algorithms/advantages.py b/relax/algorithms/advantages.py index 2c6936c1c..e3ee83b1c 100644 --- a/relax/algorithms/advantages.py +++ b/relax/algorithms/advantages.py @@ -68,13 +68,22 @@ def whiten_scalar(values: torch.Tensor, *, process_group: dist.ProcessGroup | No def _as_reward_tensor(rewards: Any, kl: list[torch.Tensor]) -> torch.Tensor: """Rewards as a detached float32 tensor on the KL tensors' device. - ``detach()`` is load-bearing, not defensive. The pre-registry code built - this with ``torch.tensor(rewards, ...)``, which copies and drops autograd - history even when handed a tensor. ``.to()`` returns the *same* object when - dtype and device already match, so without the detach a caller passing a - reward tensor with ``requires_grad=True`` would get advantages that carry - grad history into the policy loss -- a path that did not exist before and - that no caller asks for. Rewards are data, not something to backprop into. + ``detach()`` is defensive, and worth being precise about because the + comment here used to call it load-bearing, which overstates it. + + Both production call sites pass ``list[float]`` -- ``loss.py`` reads + ``rollout_data["rewards"]`` and the Advantages deployment reads the same + column off the TransferQueue -- so the tensor branch below is not currently + reached by anything, and removing the detach would change no observable + behaviour today. + + It stays because the refactor did quietly widen what this accepts. The + pre-registry code built the tensor with ``torch.tensor(rewards, ...)``, + which copies and drops autograd history even when handed a tensor; + ``.to()`` returns the *same* object when dtype and device already match. So + a future caller passing a reward tensor with ``requires_grad=True`` would, + without the detach, get advantages carrying grad history into the policy + loss. Rewards are data, not something to backprop into. """ if isinstance(rewards, torch.Tensor): return rewards.detach().to(dtype=torch.float32, device=kl[0].device) @@ -171,7 +180,19 @@ def _whiten_by_segment(values, mini_batch_sizes, process_group): # on until the whole group has shared it. local_error = "" if mini_batch_sizes is None: - n_segments = 1 + # Not a fallback. Whitening the merged rollout as one window is a + # different optimisation target, not a coarser version of the same one: + # `test_merging_the_batches_would_flip_signs_not_just_rescale` shows + # half the advantages changing sign. Every Megatron producer writes + # these counts today, so a missing one means a new or changed caller, + # and the only thing worse than that caller failing is that caller + # training on the wrong objective while every metric stays finite. + n_segments = 0 + local_error = ( + "mini_batch_sizes is None: the caller did not supply " + "`rollout_mini_local_sample_counts`. GDPO whitens per training batch, so there is no " + "safe default -- whitening the merged rollout instead optimises a different objective." + ) elif not mini_batch_sizes or any(not isinstance(n, int) or n <= 0 for n in mini_batch_sizes): n_segments = 0 local_error = f"mini_batch_sizes must be a non-empty list of positive ints, got {mini_batch_sizes}." @@ -186,7 +207,7 @@ def _whiten_by_segment(values, mini_batch_sizes, process_group): _agree_on_segmentation(n_segments, local_error, values, process_group) - if mini_batch_sizes is None or n_segments == 1: + if n_segments == 1: return whiten_scalar(values, process_group=process_group) out, start = [], 0 for size in mini_batch_sizes: diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index 04da50be2..206787a57 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -30,6 +30,14 @@ """Largest finite float32. Rewards are carried as float32 from here on, so a value above this is an overflow waiting to happen rather than a large reward.""" +_NOISE_FLOOR_SAFETY = 8.0 +"""Headroom on :func:`component_noise_scale`'s first-order error bound. + +The bound is an estimate, not a proof, so it gets an order of magnitude. It can +afford to: on well-conditioned rewards the bound sits fifteen orders of +magnitude below the signal, so any value in this neighbourhood suppresses +exactly the same things.""" + def group_positions(samples: list[Any], expected_size: int) -> dict[int, list[int]]: """Map ``Sample.group_index`` to the positions it occupies in ``samples``. @@ -212,13 +220,23 @@ def extract_reward_components(samples: list[Any], keys: list[str]) -> torch.Tens row.append(numeric) rows.append(row) - components = torch.tensor(rows, dtype=torch.float32) + # float64, not float32. The rewards arrive as Python floats and the + # standardisation that follows is ill-conditioned exactly when two + # components cancel: `C - r` loses the low bits on a float32 cast, and + # dividing what is left by a near-zero std turns those lost bits into a + # full-scale advantage. Keeping the width here is what makes the residue + # small enough for `combine_group`'s noise floor to be able to tell signal + # from rounding at all -- in float32 the residue is larger than the signal. + # The float32 range check above still stands: the combined reward is cast + # back down further along, so a value that cannot survive that cast is + # still a bug worth reporting at the reward function that produced it. + components = torch.tensor(rows, dtype=torch.float64) # Belt and braces: the per-value check above is exact, but it only sees what # `float()` returned. Anything that slips past it must not reach the # normaliser, where non-finite input is indistinguishable from a collapse. if not torch.isfinite(components).all(): bad = (~torch.isfinite(components)).nonzero()[0].tolist() - raise ValueError(f"Reward {keys[bad[1]]!r} of sample {bad[0]} is not finite after casting to float32.") + raise ValueError(f"Reward {keys[bad[1]]!r} of sample {bad[0]} is not finite as a tensor.") return components @@ -233,24 +251,12 @@ def standardize_group_components(group: torch.Tensor) -> torch.Tensor: :func:`relax.algorithms.numerics.distributed_mean_std` does: the means and stds are the part where cancellation bites, and widening them is cheap. - It does **not** fix the related failure worth knowing about. Two columns - that sum to a constant should standardise to exact opposites and cancel, - and in float64 they do -- exactly. But the columns arrive already cast to - float32 by :func:`extract_reward_components`, and a value like - ``C - r`` does not survive that cast unchanged. The pair still *sums* to - ``C`` in float32 (the addition rounds back), so the group looks flat, while - the individual values have drifted enough to leave a residue of order 1e-4 - after standardisation. Step 3 then divides that residue by a batch std of - the same order and returns advantages of order 1: a group with no signal - gets a confident gradient whose direction is decided by rounding. - - Measured on rewards summing to 308.95172119140625, the batch comes out as - [-0.4344, 0.5251, -0.0908]. GRPO does not have this failure, because it - standardises the sum, which *is* exactly constant, and its collapse check - catches it. Raising the precision here does not help -- the information was - lost before this function saw it. Documented in examples/gdpo/README.md - under the deviations; fixing it needs either a wider reward dtype end to - end or a noise floor in step 3, neither of which belongs in this PR. + Standardising is ill-conditioned when a column barely varies: it divides + by a std that is near zero, so the *relative* error in the result grows + like ``max|x| / std``. That is not a corner case here -- it is precisely + what two columns summing to a constant look like, and it is why + :func:`component_noise_scale` exists and why the columns now arrive in + float64. See :func:`combine_group` for what is done with the estimate. """ work = group.double() centered = work - work.mean(dim=0, keepdim=True) @@ -261,6 +267,78 @@ def standardize_group_components(group: torch.Tensor) -> torch.Tensor: return scaled.to(group.dtype) +def component_noise_scale(group: torch.Tensor) -> torch.Tensor: + """Per-column bound on how much of :func:`standardize_group_components` is + rounding. + + Returns one value per column: the magnitude below which that column's + standardised values are indistinguishable from the arithmetic that + produced them. + + The bound is the condition number of the standardisation. Each input + carries a relative representation error of about ``eps``, i.e. an absolute + error of ``eps * max|x|``. Dividing by ``std + GDPO_EPS`` scales that error + up by the same factor it scales the signal, so the error in the output is + ``eps * max|x| / (std + GDPO_EPS)``. + + Two properties make this usable rather than a tuning knob: + + * On a well-conditioned column it is negligible. Rewards around 1.0 with a + std of 0.1 give ``2.2e-16 * 1 / 0.1 ~ 2e-15`` against signal of order 1 + -- fifteen orders of margin, so nothing real is ever suppressed. + * On the pathological column it matches what actually happens. For the + constant-sum group in ``test_gdpo.py`` the bound comes to ~5.1e-10 and + the residue measures 4.1e-10. + + A collapsed column contributes exactly zero (not ``0/eps`` noise), so it + contributes no error either. + """ + work = group.double() + eps = torch.finfo(work.dtype).eps + std = work.std(dim=0) + scale = work.abs().amax(dim=0) + noise = eps * scale / (std + GDPO_EPS) + return torch.where(collapsed_columns(group, dim=0), torch.zeros_like(noise), noise) + + +def combine_group(group: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """GDPO steps 1 and 2 for one ``[G, K]`` prompt group, with a noise floor. + + Two components summing to a constant should standardise to exact opposites + and cancel. What is left instead is the rounding error of an + ill-conditioned division, and it does not stay small: step 3 divides by the + std of that very residue. + + Two separate things keep that from becoming a gradient, and it is worth + being exact about which does what, because the tempting summary -- "the + floor stops a full-scale fake advantage" -- is not true: + + * **The float64 columns do the heavy lifting.** In float32 the residue was + of order 1e-1, larger than the signal itself, and step 3 returned + advantages of order 1. In float64 it is of order 1e-10, and ``GDPO_EPS`` + in step 3's denominator caps the amplification, so the worst case is + already down to ~1e-6. No floor is needed to avoid a fake gradient. + * **The floor makes the cancellation exactly detectable.** 1e-6 is not + zero, and several consumers test for zero: + :func:`group_carries_reward_signal` would keep a group contributing + nothing, and the all-groups-silent warning would never fire. Rounding a + residue to the zero it mathematically is turns a "too small to matter" + into a "recognisably absent". + + The comparison is per group and all-or-nothing, matching how a collapsed + column is already handled. ``_NOISE_FLOOR_SAFETY`` is the one judgement + call; given the fifteen orders of margin on well-conditioned input, its + exact value changes nothing that is not already noise. + """ + standardized = standardize_group_components(group) + combined = (standardized.double() * weights.double()).sum(dim=1) + + floor = _NOISE_FLOOR_SAFETY * float((weights.double().abs() * component_noise_scale(group)).sum()) + if float(combined.abs().amax()) <= floor: + return torch.zeros_like(combined) + return combined + + def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: """GDPO steps 1 and 2 (arXiv 2601.05242, Eq. 4 and Eq. 7). @@ -299,28 +377,6 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl components = extract_reward_components(samples, keys) positions_by_group = group_positions(samples, args.n_samples_per_prompt) - normalized = torch.zeros_like(components) - fully_collapsed_groups = 0 - for positions in positions_by_group.values(): - group = components[positions] - normalized[positions] = standardize_group_components(group) - if bool(collapsed_columns(group, dim=0).all()): - fully_collapsed_groups += 1 - - if fully_collapsed_groups == len(positions_by_group): - # Every component collapsed in every group, so this batch produces no - # gradient at all. Usually the reward function is constant for these - # prompts (e.g. a format reward when nothing in the prompt asks for a - # format). Worth one line, because the symptom downstream is simply - # "loss does not move". - logger.warning( - "GDPO: all reward components collapsed in all %d groups of this batch (keys=%s); " - "the batch contributes no gradient. Check that each of these rewards actually varies " - "across rollouts of the same prompt.", - fully_collapsed_groups, - keys, - ) - weight_tensor = torch.tensor(weights, dtype=torch.float32) if not torch.isfinite(weight_tensor).all(): # `resolve_gdpo_weights` checked `math.isfinite` on the Python floats. @@ -339,16 +395,41 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl "the combined advantage would be identically 0." ) - combined = (normalized * weight_tensor).sum(dim=1) + combined = torch.zeros(len(samples), dtype=torch.float64) + silent_groups = 0 + for positions in positions_by_group.values(): + combined[positions] = combine_group(components[positions], weight_tensor) + if not bool(combined[positions].any()): + silent_groups += 1 + + if silent_groups == len(positions_by_group): + # Every group came out at exactly zero, so this batch produces no + # gradient at all. Usually the reward function is constant for these + # prompts (e.g. a format reward when nothing in the prompt asks for a + # format); the other way in is components that cancel. Worth one line, + # because the symptom downstream is simply "loss does not move". + logger.warning( + "GDPO: all %d groups of this batch combined to exactly zero (keys=%s); the batch " + "contributes no gradient. Either these rewards do not vary across rollouts of the " + "same prompt, or they cancel under the configured weights (weights=%s).", + silent_groups, + keys, + weights, + ) + if not torch.isfinite(combined).all(): # Every individual reward fit in float32 (checked in # `extract_reward_components`), but the arithmetic between them need - # not: a group of [3e38, 2e38, 1e38, 0] overflows while computing its - # own mean. Left unchecked the `inf` reaches `whiten_scalar`, reads as - # a non-finite std, and the batch is silently zeroed -- the exact - # failure the per-value check was added to prevent, one stage later. + # not. The mean itself no longer overflows -- the columns and the + # standardisation are float64 now, so [3e38, 2e38, 1e38, 0] averages + # fine, and the earlier version of this comment claiming otherwise is + # out of date. What can still reach infinity is the weighting: float64 + # weights near its own maximum, or a reward near float64's range rather + # than float32's. Left unchecked the `inf` reaches `whiten_scalar`, + # reads as a non-finite std, and the batch is silently zeroed -- the + # exact failure the per-value check exists to prevent, one stage later. raise ValueError( - "GDPO produced a non-finite combined advantage from finite inputs; the float32 " + "GDPO produced a non-finite combined advantage from finite inputs; the " f"arithmetic overflowed. Rescale the rewards (keys={keys}) or their weights." ) return combined.tolist() @@ -405,8 +486,80 @@ def group_carries_reward_signal(args: Any, samples: list[Any]) -> bool: keys = resolve_gdpo_keys(args) weights = torch.tensor(resolve_gdpo_weights(args, keys), dtype=torch.float32) components = extract_reward_components(samples, keys) - combined = (standardize_group_components(components) * weights).sum(dim=1) - return bool((combined != 0).any()) + # `combine_group`, not steps 1 and 2 open-coded: this has to answer the + # same question `normalize_gdpo_decoupled` will, including the noise floor. + # Open-coding it is how the two came apart before. + return bool(combine_group(components, weights).any()) + + +def observed_reward_signal(args: Any, samples: list[Any]) -> bool | None: + """:func:`group_carries_reward_signal` for callers that only report. + + Returns ``None`` -- "cannot tell" -- where the strict version raises. + + This is not the strict check with its errors swallowed; it is a different + question, and the difference is not cosmetic. The strict version is asked + by the dynamic-sampling filter, which decides whether a group enters + *training*, so a reward it cannot read is a bug that should stop the run. + The zero-std metrics are asked by the logger, and two things follow: + + * A metric must not decide whether training proceeds. Making observation + the first enforcer of a contract means a broken reward function is + reported at the log line rather than at the stage that consumes it, and + it takes the rollout down on the way. + * The contract is not even in force everywhere the metrics run. Eval may + use a different reward model (``EvalConfig.rm_type``), so an eval reward + legitimately need not carry the ``--gdpo-reward-keys`` that training + needs -- nothing in eval consumes them. The strict question has no + correct answer there; ``None`` is the honest one. + + During training the same violation is still raised, by + :func:`normalize_gdpo_decoupled`, which is the stage that actually reads + the components. Nothing is hidden -- but it is logged here too, because a + group the metrics could not read is worth knowing about either way. + """ + try: + return group_carries_reward_signal(args, samples) + except (TypeError, ValueError) as exc: + logger.warning( + "zero-std metrics: skipping a group whose reward could not be read (%s). " + "For a training rollout the reward stage will raise on this; for eval it may just " + "mean the eval reward model returns a different schema, which is fine.", + exc, + ) + return None + + +def metrics_group_verdict(args: Any, samples: list[Any]) -> bool | None: + """Is this prompt group flat, for zero-std reporting? ``None`` if + unanswerable. + + ``True`` flat, ``False`` varying, ``None`` neither -- nothing in the group + was scored, or the reward could not be read. + + This lives here rather than in the two rollout modules because there are + *two* copies of ``_compute_zero_std_metrics``, one in + ``relax/agentic/rollout.py`` and one in ``relax/distributed/ray/rollout.py``, + and they have already drifted apart once: the agentic one dropped unscored + samples and the distributed one did not, which turned a `reward=None` into + a ``TypeError`` on the rollout's way out. Only the agentic one is reachable + from a CPU test -- the other pulls in ``sglang`` at import -- so a shared + helper is the only version of this logic that can be tested at all. + + ``reward=None`` is a real state, not a defensive check: under ``--group-rm`` + the group reward is assigned in one shot that is skipped entirely when the + rollout aborts, which is why the "reward is not None" assert in + ``sglang_rollout.py`` exempts ``group_rm`` in the first place. + + Callers still differ on what to do with ``None``, and that difference is + deliberate -- each preserves what its own metric reported before. See the + call sites. + """ + rewarded = [sample for sample in samples if sample.reward is not None] + if not rewarded: + return None + signal = observed_reward_signal(args, rewarded) + return None if signal is None else not signal REWARD_NORMALIZERS: dict[str, Callable[[Any, list[Any], list[float]], list[float]]] = { diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index f96462cb6..03af0d424 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -591,10 +591,11 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) # rollout_data is this rank's shard of the WHOLE rollout, not of one # training batch: actor.py collects `num_rollout_minis` windows of # global_batch_size/dp_size and concat_rollout_batches merges them before - # this call ("we may need normalize the whole rollout", actor.py). So the - # reduction below makes the statistic describe the rollout across the DP - # group -- which spans every optimizer step in it, not one batch. See the - # known-deviation note in docs/*/examples/algorithms.md. + # this call ("we may need normalize the whole rollout", actor.py). This + # group is therefore the DP group across that merged rollout -- it says + # *who* to reduce with, and `mini_batch_sizes` below says *where the + # batch boundaries are* inside it. Both are needed: the group alone + # would make one statistic span every optimizer step in the rollout. process_group=mpu.get_data_parallel_group(), # Per-training-batch counts, written by actor.py before it merges the # rollout. Passing them is what lets a batch-level statistic describe one diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index 776f10211..1f46d02e4 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -22,7 +22,7 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS from relax.algorithms import get_algorithm -from relax.algorithms.rewards import group_carries_reward_signal +from relax.algorithms.rewards import metrics_group_verdict from relax.backends.sglang.sglang_engine import SGLangEngine from relax.distributed.ray.rollout_validation import validate_server_group_gpu_indices from relax.engine.rollout.base_types import call_rollout_fn @@ -4057,12 +4057,13 @@ def _compute_zero_std_metrics(args, all_samples: list[Sample]): return {} def _is_zero_std(samples: list[Sample]): - # Reads whichever notion of "signal" this algorithm uses: the - # --reward-key scalar for single-reward algorithms, every component for - # multi-reward ones. Counting a GDPO group as zero-std because its - # summed reward is flat overstates the count and reads as "most of the - # batch is dead" when it is not. - return not group_carries_reward_signal(args, samples) + # `is not False`, so an unreadable or entirely unscored group counts as + # zero-std. That is what this metric reported before: its predicate was + # `len(rewards) == 0 or all(rewards[0] == r ...)`, and an all-None group + # satisfied the `all(...)`. The agentic copy skips such groups instead, + # because *its* predicate dropped them before counting. Neither is more + # correct; each keeps its own metric comparable across this change. + return metrics_group_verdict(args, samples) is not False all_sample_groups = group_by(all_samples, lambda s: s.group_index) interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)] diff --git a/relax/engine/filters/dynamic_sampling_filters.py b/relax/engine/filters/dynamic_sampling_filters.py index 5fe387b90..be816fbf3 100644 --- a/relax/engine/filters/dynamic_sampling_filters.py +++ b/relax/engine/filters/dynamic_sampling_filters.py @@ -14,10 +14,16 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): "No signal" is algorithm-dependent, which is why the test is not spelled out here. For a single-reward algorithm it is whether the ``--reward-key`` scalar varies in float32 -- equivalent to the ``std > 0`` this replaced, - since ``min == max`` and ``std == 0`` agree on every float32 input. For a - multi-reward algorithm it is whether *every* component is flat; judging - those by the summed scalar would drop the groups the algorithm exists to - keep. + since ``min == max`` and ``std == 0`` agree on every float32 input. + + For a multi-reward algorithm it is whether the *combined* advantage comes + out non-zero -- not whether every component is flat, which is a weaker + condition and was what this docstring used to claim. The two part company + exactly where it matters: a zero weight mutes a varying component, and two + components whose standardised values are opposites cancel. Both leave a + group that varies component-wise and still contributes no gradient. Judging + by the summed ``--reward-key`` scalar is wrong in the other direction, and + drops groups the algorithm exists to keep. Note this can now raise rather than merely returning a verdict: for a multi-reward algorithm it runs the same component extraction the reward diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index a0f8c11e6..9f30c8c75 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -2,6 +2,7 @@ import argparse import json +import math import os import sys import warnings @@ -3104,6 +3105,26 @@ def _validate_multi_reward_args(args, spec) -> None: f"`--gdpo-reward-weights` has {len(weights)} entries but `--gdpo-reward-keys` has {len(keys)}." ) + # The same two conditions the reward stage rejects, asked here instead of at + # the first rollout. `type=float` accepts "nan" and "inf" from the command + # line, and 1e-50 is a perfectly good Python float that flushes to zero once + # it reaches float32. Neither depends on any rollout data, so waiting until + # one exists only means the cluster is up and the checkpoint is loaded + # before the typo is reported. The reward-stage checks stay: they also cover + # weights that arrive from a config override after this point. + if weights is not None: + import torch + + unusable = [w for w in weights if not math.isfinite(w)] + if unusable: + raise ValueError(f"`--gdpo-reward-weights` contains non-finite values: {unusable}.") + survives_cast = torch.tensor(weights, dtype=torch.float32) + if not bool(survives_cast.abs().sum() > 0): + raise ValueError( + f"`--gdpo-reward-weights` {weights} are all zero in float32; the combined advantage " + "would be identically 0 and the run would train on no signal." + ) + # Components arrive as a dict; without --reward-key the raw_reward column # would hold dicts, which the TransferQueue conversion cannot represent. if not args.reward_key: diff --git a/tests/algorithms/test_advantage_estimators.py b/tests/algorithms/test_advantage_estimators.py index 0b07feace..dab86d910 100644 --- a/tests/algorithms/test_advantage_estimators.py +++ b/tests/algorithms/test_advantage_estimators.py @@ -29,6 +29,11 @@ def _inputs(lengths=(3, 2)): response_lengths=list(lengths), total_lengths=[n + 2 for n in lengths], values=None, + # One segment covering the whole shard. Estimators that do not whiten + # per batch absorb this in `**_unused`; GDPO requires it, because a + # caller that cannot say how the rollout was split cannot be given a + # default without silently changing which objective it optimises. + mini_batch_sizes=[len(lengths)], ) diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py index eddbaefa6..80bea5acd 100644 --- a/tests/algorithms/test_arguments_spec_driven.py +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -476,7 +476,13 @@ def test_supports_fully_async_is_declared_wherever_it_is_false(): def test_dynamic_sampling_filter_warns_for_multi_reward(arguments_module, caplog): - """The built-in filter judges a group by the single --reward-key scalar.""" + """A *custom* filter is opaque, so a multi-reward run warns about it. + + Not the built-in one, which this docstring used to describe: that reads + `group_carries_reward_signal` and judges the group by the combined + advantage. The warning exists because a custom filter may still reduce the + group to the single --reward-key scalar, and from here we cannot tell. + """ import logging with caplog.at_level(logging.WARNING): @@ -543,3 +549,33 @@ def test_a_typo_in_advantage_normalization_is_rejected(arguments_module): broken = dataclasses.replace(get_algorithm("grpo"), advantage_normalization="typo") with pytest.raises(ValueError, match="advantage_normalization"): arguments_module._assert_spec_implementations_resolve(broken) + + +# ---------------- weights are checked before the cluster starts ---------------- + + +@pytest.mark.parametrize( + "weights, expected", + [ + ([float("nan"), 1.0], "non-finite"), + ([float("inf"), 1.0], "non-finite"), + ([0.0, 0.0], "all zero in float32"), + ([1e-50, 1e-50], "all zero in float32"), + ], +) +def test_unusable_gdpo_weights_are_rejected_at_startup(arguments_module, weights, expected): + """`--gdpo-reward-weights nan 1` used to start the whole run first. + + `type=float` accepts "nan" and "inf" from a shell, and 1e-50 is a real + Python float that flushes to zero in float32. None of it depends on rollout + data, so failing at the first rollout only means the placement groups are + up and the checkpoint is loaded before the typo is reported. + """ + args = _args(gdpo_reward_weights=weights) + with pytest.raises(ValueError, match=expected): + arguments_module.validate_algorithm_args(args) + + +def test_usable_gdpo_weights_still_pass(arguments_module): + """A weight may legitimately be zero, as long as they are not all zero.""" + arguments_module.validate_algorithm_args(_args(gdpo_reward_weights=[0.0, 1.0])) diff --git a/tests/algorithms/test_dispatch_parity_vs_main.py b/tests/algorithms/test_dispatch_parity_vs_main.py index d8deb27d5..7d762c711 100644 --- a/tests/algorithms/test_dispatch_parity_vs_main.py +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -556,3 +556,139 @@ def test_loss_py_actually_forwards_the_mini_batch_boundaries(): "loss.py must pass the per-training-batch counts; without them GDPO's step 3 " "normalises over the whole merged rollout again" ) + + +# ---------------- advantage_normalization ---------------- +# +# This field is read on every optimizer step of every algorithm, and until +# these tests existed nothing pinned its value: the only test naming it +# rejected an invalid string. Dropping `"token_global"` from a REINFORCE++ spec +# (silently reverting it to the `"whiten"` default) or flipping the comparison +# in loss.py left the whole suite green. + + +@pytest.mark.parametrize("name", MAIN_ALGORITHMS) +def test_advantage_normalization_matches_main(name): + expected = "token_global" if name in MAIN_TOKEN_GLOBAL else "whiten" + assert get_algorithm(name).advantage_normalization == expected + + +def test_exactly_mains_algorithms_take_the_token_global_path(): + """Pins the set, not just the members. + + A per-algorithm check cannot fail when a *new* algorithm defaults into the + wrong branch, and the default is the branch every existing algorithm but + two takes -- so the mistake it would miss is the likely one. + """ + token_global = {n for n in list_algorithm_names() if get_algorithm(n).advantage_normalization == "token_global"} + assert token_global == MAIN_TOKEN_GLOBAL + + +def test_loss_py_selects_both_behaviours_from_the_field(): + """Both of main's call sites must read the spec, not the algorithm name. + + Source-level because importing `loss.py` needs Megatron. It is paired with + the value tests above: those pin what the field says, this pins that the + two branches still ask it. + """ + import pathlib + + import relax + + src = (pathlib.Path(relax.__file__).parent / "backends/megatron/loss.py").read_text() + + assert src.count('advantage_normalization == "token_global"') == 2 + for literal in MAIN_TOKEN_GLOBAL: + assert f'"{literal}"' not in src, f"loss.py still names {literal} directly" + + +def test_gae_adapter_reproduces_mains_reward_shaping(cp_disabled): + """PPO's adapter, run against a transcription of main's inline branch. + + The only estimator adapter with no numerical check: `advantage_gae` had a + co_names test, which cannot see a dropped argument, a sign flip on the KL + coefficient, or the terminal-reward injection landing on the wrong token. + PPO is also the algorithm whose surroundings this refactor changed most + (it is the one with a critic) and the one with no GPU smoke, so a bug here + would have had the least chance of being caught anywhere else. + + main @ 98a1274 loss.py:584-602, transcribed rather than imported: + + 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, + ) + """ + from relax.algorithms.advantages import compute_advantages_and_returns + + def _inputs(): + return dict( + rewards=[1.5, -2.0], + kl=[torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.4, 0.5])], + values=[torch.tensor([0.7, 0.8, 0.9]), torch.tensor([1.1, 1.2])], + response_lengths=[3, 2], + total_lengths=[5, 4], + ) + + args = _args("ppo", kl_coef=0.3, gamma=0.95, lambd=0.9) + + # main's branch, on its own copy of the tensors -- the shaping mutates `kl` + # in place (`k *= kl_coef`), so the two runs must not share them. + ref = _inputs() + shaped = [] + for reward, k in zip(ref["rewards"], ref["kl"], strict=False): + k *= -args.kl_coef + k[-1] += reward # cp_rank == 0 under `cp_disabled` + shaped.append(k) + want_adv, want_ret = ppo_utils.get_advantages_and_returns_batch( + ref["total_lengths"], + ref["response_lengths"], + ref["values"], + shaped, + args.gamma, + args.lambd, + padded_total_lengths=None, + ) + + got_adv, got_ret = compute_advantages_and_returns(args, **_inputs()) + + for left, right in zip(got_adv, want_adv, strict=True): + assert torch.equal(left, right) + for left, right in zip(got_ret, want_ret, strict=True): + assert torch.equal(left, right) + + +def test_gae_adapter_actually_uses_kl_coef_and_gamma(cp_disabled): + """Guards the transcription above from agreeing by coincidence. + + If the adapter ignored either argument, the test above would still pass as + long as the reference ignored it too. + """ + from relax.algorithms.advantages import compute_advantages_and_returns + + def _run(**overrides): + knobs = dict(kl_coef=0.3, gamma=0.95, lambd=0.9) + knobs.update(overrides) + adv, _ = compute_advantages_and_returns( + _args("ppo", **knobs), + rewards=[1.5, -2.0], + kl=[torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.4, 0.5])], + values=[torch.tensor([0.7, 0.8, 0.9]), torch.tensor([1.1, 1.2])], + response_lengths=[3, 2], + total_lengths=[5, 4], + ) + return torch.cat(adv) + + baseline = _run() + assert not torch.allclose(baseline, _run(kl_coef=0.9)) + assert not torch.allclose(baseline, _run(gamma=0.5)) + assert not torch.allclose(baseline, _run(lambd=0.1)) diff --git a/tests/algorithms/test_distributed_whitening.py b/tests/algorithms/test_distributed_whitening.py index 3137c3bef..bab244951 100644 --- a/tests/algorithms/test_distributed_whitening.py +++ b/tests/algorithms/test_distributed_whitening.py @@ -200,11 +200,22 @@ def test_mismatched_segment_counts_fail_on_every_rank(): assert "same number of segments" in out[rank], out[rank] -def test_no_segmentation_on_one_rank_is_also_a_mismatch(): +def test_no_segmentation_on_one_rank_fails_both_ranks(): + """Absent counts are now rejected outright, and still fail the group as + one. + + This used to surface as a segment-count mismatch (one rank planning a + single window against the other's two). Since `None` became an error in its + own right, rank 0 reports that directly -- but the property this test is + really for is unchanged and is the reason it cannot be a single-rank test: + rank 0 must not leave the collectives on its own while rank 1 waits inside + them, so rank 1 has to fail too. + """ out = _spawn("none_versus_segmented") for rank in (0, 1): assert isinstance(out[rank], str), f"rank {rank} did not raise: {out[rank]!r}" - assert "same number of segments" in out[rank], out[rank] + assert "mini_batch_sizes is None" in out[0], out[0] + assert "another rank reported malformed" in out[1], out[1] def test_malformed_metadata_on_one_rank_fails_both(): diff --git a/tests/algorithms/test_gdpo.py b/tests/algorithms/test_gdpo.py index de5dbea2a..5348873f1 100644 --- a/tests/algorithms/test_gdpo.py +++ b/tests/algorithms/test_gdpo.py @@ -17,7 +17,17 @@ from relax.algorithms import get_algorithm # noqa: E402 from relax.algorithms.advantages import whiten_scalar # noqa: E402 -from relax.algorithms.rewards import REWARD_NORMALIZERS, extract_reward_components # noqa: E402 +from relax.algorithms.rewards import ( # noqa: E402 + REWARD_NORMALIZERS, + combine_group, + component_noise_scale, + extract_reward_components, + group_carries_reward_signal, + standardize_group_components, +) + + +_UNSET = object() def _args(keys=("correctness", "format"), weights=None, n=4): @@ -185,7 +195,36 @@ def test_default_weights_are_all_ones(): assert _normalize(_args(weights=None), samples) == _normalize(_args(weights=[1.0, 1.0]), samples) -def test_grouping_uses_group_index(): +def test_grouping_follows_group_index_not_position(): + """Interleave two groups so position and group_index disagree. + + The previous version of this test only ever built the contiguous layout + [0,0,0,0,1,1,1,1], which an implementation that chopped the batch into + fixed-size runs would satisfy just as well -- it asserted nothing its own + name claimed. Here sample i belongs to group i % 2, so any position-based + grouping produces different numbers. + """ + groups = [0, 1, 0, 1, 0, 1, 0, 1] + correctness = [1.0, 5.0, 0.0, 5.0, 1.0, 5.0, 0.0, 5.0] + fmt = [1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0] + + interleaved = _normalize(_args(n=4), _mk(groups, correctness, fmt)) + + # Group 1 (odd positions) has collapsed correctness, so only format acts there. + odd = [interleaved[i] for i in (1, 3, 5, 7)] + assert abs(sum(odd)) < 1e-5 + # Reordering the samples so each group is contiguous must not change any + # sample's own advantage -- which is only true if group_index decides. + order = [0, 2, 4, 6, 1, 3, 5, 7] + contiguous = _normalize( + _args(n=4), + _mk([groups[i] for i in order], [correctness[i] for i in order], [fmt[i] for i in order]), + ) + for new_position, original in enumerate(order): + assert abs(contiguous[new_position] - interleaved[original]) < 1e-5 + + +def test_collapsed_component_in_one_group_only(): correctness = [1.0, 0.0, 1.0, 0.0, 5.0, 5.0, 5.0, 5.0] fmt = [1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0] contiguous = _normalize(_args(), _mk([0, 0, 0, 0, 1, 1, 1, 1], correctness, fmt)) @@ -369,7 +408,7 @@ def test_extract_reward_components_shape_and_dtype(): samples = _mk([0, 0, 0, 0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]) out = extract_reward_components(samples, ["correctness", "format"]) assert out.shape == (4, 2) - assert out.dtype == torch.float32 + assert out.dtype == torch.float64 def test_extract_reward_components_preserves_key_order(): @@ -382,7 +421,7 @@ def test_extract_reward_components_preserves_key_order(): def test_extract_accepts_python_ints(): samples = [_S(0, {"correctness": 1, "format": 0}) for _ in range(2)] out = extract_reward_components(samples, ["correctness", "format"]) - assert out.dtype == torch.float32 + assert out.dtype == torch.float64 def test_single_reward_gdpo_is_a_constant_positive_multiple_of_grpo(): @@ -424,7 +463,7 @@ def test_warns_once_when_every_component_collapses_in_every_group(caplog): out = _normalize(_args(), samples) assert out == [0.0] * 4 - assert sum("all reward components collapsed" in r.message for r in caplog.records) == 1 + assert sum("combined to exactly zero" in r.message for r in caplog.records) == 1 def test_does_not_warn_when_some_signal_survives(caplog): @@ -434,7 +473,7 @@ def test_does_not_warn_when_some_signal_survives(caplog): with caplog.at_level(logging.WARNING): _normalize(_args(), samples) - assert not any("all reward components collapsed" in r.message for r in caplog.records) + assert not any("combined to exactly zero" in r.message for r in caplog.records) def test_does_not_warn_when_only_some_groups_collapse(caplog): @@ -448,7 +487,7 @@ def test_does_not_warn_when_only_some_groups_collapse(caplog): with caplog.at_level(logging.WARNING): _normalize(_args(), samples) - assert not any("all reward components collapsed" in r.message for r in caplog.records) + assert not any("combined to exactly zero" in r.message for r in caplog.records) # ---------------- reward value contract ---------------- @@ -465,7 +504,7 @@ def test_numpy_scalars_are_accepted(): _S(0, {"correctness": dtype(0), "format": dtype(1)}), ] out = extract_reward_components(samples, ["correctness", "format"]) - assert out.dtype == torch.float32, dtype + assert out.dtype == torch.float64, dtype assert out.shape == (2, 2), dtype @@ -579,8 +618,10 @@ def test_whitening_scope_is_whatever_the_caller_passes_not_a_training_batch(): statistics describe one training batch "matching the paper", and they do not. The Megatron caller merges `num_rollout_minis` windows before calling (actor.py: concat_rollout_batches, under the comment "we may need normalize - the whole rollout"), so with the shipped example's 4 * 8 against a - global_batch_size of 16 one whitening spans two optimizer steps. + the whole rollout"), so whether one whitening spans one optimizer step or + several is the caller's business, not this function's. The shipped example + runs 4 * 8 against a --global-batch-size of 32, i.e. one window; halving + that to 16 gives two, which is what `mini_batch_sizes` then separates. Concretely: whitening two batches together is not the same as whitening each. """ @@ -604,9 +645,17 @@ def test_whitening_scope_is_whatever_the_caller_passes_not_a_training_batch(): # ---------------- step 3 now normalises per training batch ---------------- -def _gdpo_adv(rewards, mini_batch_sizes=None): +def _gdpo_adv(rewards, mini_batch_sizes=_UNSET): + """`mini_batch_sizes` defaults to one segment covering everything. + + Not `None`: that is now rejected outright, because a caller that fails to + say how the rollout was split is asking for a different objective without + knowing it. Tests that want the merged behaviour ask for it explicitly. + """ from relax.algorithms.advantages import compute_advantages_and_returns + if mini_batch_sizes is _UNSET: + mini_batch_sizes = [len(rewards)] kl = [torch.zeros(1) for _ in rewards] adv, _ = compute_advantages_and_returns( SimpleNamespace(advantage_estimator="gdpo", kl_coef=0.0), @@ -641,7 +690,10 @@ def test_merging_the_batches_would_flip_signs_not_just_rescale(): """ first, second = [0.9, 1.1, 0.8, 1.2], [-1.2, -0.8, -1.1, -0.9] per_batch = _gdpo_adv(first + second, mini_batch_sizes=[4, 4]) - merged = _gdpo_adv(first + second, mini_batch_sizes=None) + # `[8]`, not `None`: one segment spanning the merged rollout is exactly the + # behaviour a missing count would have produced, and it is now the only way + # to ask for it -- which is the point of this test. + merged = _gdpo_adv(first + second, mini_batch_sizes=[8]) assert int((per_batch.sign() != merged.sign()).sum()) == 4 @@ -649,7 +701,19 @@ def test_merging_the_batches_would_flip_signs_not_just_rescale(): def test_a_single_batch_is_unchanged_by_the_counts(): """num_rollout_minis == 1 is the common case and must not move.""" rewards = [1.0, 2.0, 3.0, 4.0] - torch.testing.assert_close(_gdpo_adv(rewards, [4]), _gdpo_adv(rewards, None)) + torch.testing.assert_close(_gdpo_adv(rewards, [4]), _gdpo_adv(rewards, [len(rewards)])) + + +def test_absent_counts_are_rejected_rather_than_merged(): + """The caller must say how the rollout was split; there is no default. + + Falling back to one window looks like a conservative default and is not: + `test_merging_the_batches_would_flip_signs_not_just_rescale` shows it + changing the sign of half the advantages. A caller that forgot the metadata + would have trained on a different objective with every metric finite. + """ + with pytest.raises(ValueError, match="mini_batch_sizes is None"): + _gdpo_adv([1.0, 2.0, 3.0, 4.0], mini_batch_sizes=None) def test_counts_that_do_not_cover_the_shard_are_rejected(): @@ -702,3 +766,92 @@ def test_a_single_segment_is_still_size_checked(): window.""" with pytest.raises(ValueError, match="sum to"): _gdpo_adv([1.0, 2.0, 3.0, 4.0, 5.0], mini_batch_sizes=[4]) + + +# ---------------- the noise floor (see combine_group) ---------------- + + +_CONSTANT_SUM_C = 308.95172119140625 +_CONSTANT_SUM_R = [-75.74329876632146, -75.74330217449115, -75.7432989215475] +"""Three samples whose two components sum to exactly the same float32 value. + +Found by search, not by hand: the pair has to round to an identical sum while +the individual values still differ, which needs the components to sit far +enough from zero that their ulp exceeds the spread being represented. +""" + + +def _constant_sum_group(): + return torch.tensor([[r, _CONSTANT_SUM_C - r] for r in _CONSTANT_SUM_R], dtype=torch.float64) + + +def test_components_that_cancel_produce_exactly_zero_not_amplified_rounding(): + """Equal weights on two components summing to a constant must cancel. + + In float32 this group came out of step 2 at [-0.086, 0.173, -0.086] and out + of step 3 -- which divides by the std of that very residue -- at [-0.577, + 1.154, -0.577]. Finite, plausible, and pointing wherever the last bits of + the reward happened to fall. + """ + combined = combine_group(_constant_sum_group(), torch.tensor([0.5, 0.5])) + assert combined.tolist() == [0.0, 0.0, 0.0] + assert whiten_scalar(combined).tolist() == [0.0, 0.0, 0.0] + + +def test_float32_columns_are_what_made_the_residue_dangerous(): + """The dtype is the half that stops a full-scale fake advantage. + + Pins the claim `combine_group` makes about the split of responsibility. If + someone narrows `extract_reward_components` back to float32 on the grounds + that "the floor handles it", the residue returns to being larger than the + signal and this fails. + """ + weights = torch.tensor([0.5, 0.5]) + narrow = _constant_sum_group().float().double() + + residue = (standardize_group_components(narrow).double() * weights.double()).sum(dim=1) + assert residue.abs().amax() > 1e-3 + assert whiten_scalar(residue).abs().amax() > 0.5 + + +def test_float64_alone_leaves_a_residue_small_but_not_zero(): + """The floor's job is exact detectability, not magnitude. + + Also pins that it is `GDPO_EPS` capping step 3, not luck: without it the + residue would be divided by its own std and come back to order 1. + """ + group = _constant_sum_group() + weights = torch.tensor([0.5, 0.5]) + unfloored = (standardize_group_components(group).double() * weights.double()).sum(dim=1) + + assert 0.0 < float(unfloored.abs().amax()) < 1e-8 + assert float(whiten_scalar(unfloored).abs().amax()) < 1e-4 + + +def test_the_floor_leaves_a_well_conditioned_group_alone(): + """The margin is fifteen orders of magnitude, so nothing real is + suppressed.""" + group = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.0, 0.0]], dtype=torch.float64) + weights = torch.tensor([0.5, 0.5]) + + combined = combine_group(group, weights) + floor = component_noise_scale(group) + + assert combined.abs().amax() > 0.1 + assert float(floor.amax()) < 1e-14 + + +def test_noise_scale_is_zero_for_a_collapsed_column(): + """A collapsed column contributes exact zeros, so it contributes no + error.""" + group = torch.tensor([[1.0, 5.0], [0.0, 5.0], [1.0, 5.0]], dtype=torch.float64) + assert component_noise_scale(group)[1].item() == 0.0 + + +def test_the_filter_agrees_with_the_normalizer_on_a_cancelling_group(): + """Both go through `combine_group`, so they cannot disagree about this.""" + samples = [_S(0, {"correctness": r, "format": _CONSTANT_SUM_C - r}) for r in _CONSTANT_SUM_R] + args = _args(weights=[0.5, 0.5], n=3) + + assert group_carries_reward_signal(args, samples) is False + assert REWARD_NORMALIZERS["gdpo_decoupled"](args, samples, [0.0] * 3) == [0.0, 0.0, 0.0] diff --git a/tests/algorithms/test_multi_reward_consumers.py b/tests/algorithms/test_multi_reward_consumers.py index 2a2badccd..e56d97391 100644 --- a/tests/algorithms/test_multi_reward_consumers.py +++ b/tests/algorithms/test_multi_reward_consumers.py @@ -27,7 +27,11 @@ torch = pytest.importorskip("torch") -from relax.algorithms.rewards import group_carries_reward_signal, normalize_gdpo_decoupled # noqa: E402 +from relax.algorithms.rewards import ( # noqa: E402 + group_carries_reward_signal, + metrics_group_verdict, + normalize_gdpo_decoupled, +) def _combined(args, group): @@ -46,6 +50,14 @@ def get_reward_value(self, args): return self.reward if not args.reward_key else self.reward[args.reward_key] def get_reward_components(self, keys): + # Mirrors Sample.get_reward_components, including the error type: a + # double that raises KeyError where the real one raises ValueError + # cannot exercise any caller that distinguishes them. + if not isinstance(self.reward, dict): + raise ValueError(f"Sample.reward must be a dict to read components {keys}") + for key in keys: + if key not in self.reward: + raise ValueError(f"Reward key {key!r} missing from sample reward") return [self.reward[key] for key in keys] @@ -207,3 +219,80 @@ def test_builtin_filter_is_unchanged_for_single_reward_algorithms(): varied = _group(correctness=[1.0, 0.0, 1.0, 0.0], fmt=[1.0, 0.0, 1.0, 0.0]) assert check_reward_nonzero_std(_grpo_args(), varied).keep is True + + +# ---------------- the zero-std metrics (metrics_group_verdict) ---------------- +# +# There are two copies of `_compute_zero_std_metrics`, and only the agentic one +# is importable without `sglang`. Both now delegate the whole decision to +# `metrics_group_verdict`, so testing it here covers the copy that cannot be +# imported. That indirection is the point: the copies drifted apart once and +# the drift was invisible precisely because neither had a test. + + +def _scalar_args(): + """GRPO reading the bare reward, not a dict key -- the common single-reward + setup.""" + return SimpleNamespace(advantage_estimator="grpo", n_samples_per_prompt=4, reward_key=None) + + +def test_metrics_verdict_tolerates_an_unscored_sample(): + """`reward=None` reaches the metrics under --group-rm when a rollout + aborts. + + The scalar path builds a float32 tensor from these; a None in it is a + TypeError, raised from a logging helper on the rollout's way out. Before + the shared helper the distributed copy did exactly that. + """ + group = [_S(0, None), _S(0, 1.0), _S(0, 1.0)] + assert metrics_group_verdict(_scalar_args(), group) is True + + +def test_metrics_verdict_is_unknown_when_nothing_was_scored(): + assert metrics_group_verdict(_scalar_args(), [_S(0, None), _S(0, None)]) is None + + +def test_metrics_verdict_still_reads_variation(): + assert metrics_group_verdict(_scalar_args(), [_S(0, 0.0), _S(0, 1.0)]) is False + + +def test_metrics_verdict_is_unknown_when_the_reward_schema_does_not_fit(): + """Eval may run a different reward model than training + (EvalConfig.rm_type). + + Its rewards need not carry --gdpo-reward-keys, and nothing in eval consumes + them. The strict question has no correct answer here, so the observer must + not be the thing that enforces the training contract -- it took eval down. + """ + args = SimpleNamespace( + advantage_estimator="gdpo", + gdpo_reward_keys=["a", "b"], + gdpo_reward_weights=None, + reward_key="score", + n_samples_per_prompt=2, + ) + group = [_S(0, {"score": 1.0, "a": 1.0}), _S(0, {"score": 2.0, "a": 2.0})] + + assert metrics_group_verdict(args, group) is None + # ...while the stage that actually consumes the components still refuses it. + with pytest.raises(ValueError, match="missing from sample reward"): + group_carries_reward_signal(args, group) + + +def test_the_two_metrics_copies_both_delegate_the_decision(): + """Neither copy may re-derive "is this group flat" for itself. + + A source check because `relax/distributed/ray/rollout.py` imports `sglang` + and cannot be loaded here. It is weak -- it cannot tell whether the verdict + is *used* correctly -- so it is paired with the behavioural tests above + rather than standing in for them. + """ + import pathlib + + import relax + + root = pathlib.Path(relax.__file__).parent + for path in ("agentic/rollout.py", "distributed/ray/rollout.py"): + src = (root / path).read_text() + assert "metrics_group_verdict" in src, path + assert "group_carries_reward_signal" not in src, f"{path} must not ask the strict question" From 4be7743e96c0562713fea2445db439d4f9028642 Mon Sep 17 00:00:00 2001 From: Ziheng Zhang Date: Fri, 14 Aug 2026 19:08:44 +0800 Subject: [PATCH 10/22] fix(rewards): stop a crashing label and a signal-swallowing floor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two defects on either side of `combine_group`, both from a decision that was made in the wrong place rather than made wrongly. # 🐛 Bug Fix ## The zero-std metrics still died on an unscored sample `metrics_group_verdict` moved the *verdict* into one testable place and left the *label* behind in both copies of `_compute_zero_std_metrics`, where they promptly disagreed: the agentic copy reads it off the first scored sample, the distributed copy off `group[0]`, scored or not. The crash it was introduced to remove did not go away, it moved one line down -- from `torch.tensor([None], dtype=torch.float32)` to `round(None, 1)`, still a `TypeError` out of a logging helper on the rollout's way out. Two ways in, and they are not the same bug: - A group where nothing was scored. The distributed copy counts those deliberately (`is not False`), then asks for a label that cannot exist. Reachable on `98a1274` too, where the old predicate also counted them: pre-existing, not a regression. - A flat group whose first sample is the unscored one. `[None, 0.5, 0.5]` verdicts as flat and `group[0]` is the None. The old predicate returned False here and never reached the label, so this one *is* a regression. `zero_std_group_label` now answers "which reward does this group get filed under" in one place, off the first scored sample, returning `None` when there is no scored sample at all -- not a third policy, just the absence of a label to file it under. Both callers drop that group. The source assertion that pinned "both copies delegate the verdict" now pins the label too, since splitting one decision out and leaving its other half behind is how these copies drifted the second time. ## The noise floor could zero a group that carries signal `component_noise_scale` is `eps * max|x| / (std + GDPO_EPS)`, so it grows with the ratio of a reward's magnitude to its own spread. The claim it shipped with -- fifteen orders of margin, nothing real is ever suppressed -- holds on well-conditioned input, which is where the only test measured it. It does not hold generally: | base | spread | floor | \|combined\|max | zeroed | | ---- | ------ | ------ | --------------- | ------ | | 1e9 | 0.1 | 1.8e-5 | 0.999 | no | | 1e12 | 1e-3 | 1.65 | 0.907 | yes | | 1e15 | 1.0 | 1.78 | 0.9999 | yes | The last two carry an order-1 combined advantage and are zeroed anyway. Such a reward is finite, well under `_FLOAT32_MAX`, and passes every check in `extract_reward_components`. Zeroing it is silent in both directions: groups below the floor are dropped by the filter with no log at all, and when every group is, the batch warning names two causes -- rewards that do not vary, weights that cancel -- neither of which is what happened, so it reads as a reward-function bug that is not there. `extract_reward_components` already refuses a reward that overflows float32 rather than casting it to `inf`, on the grounds that a silently zeroed component is indistinguishable from a genuinely collapsed one. A reward with no significant digits where it varies is the same failure one stage later, and now gets the same answer. `_MAX_COMPONENT_NOISE` is 0.01. The values it bounds have unit variance, so it reads as a fraction: at 1% of the standardised value being rounding, the gradient's direction is partly noise whether or not the floor zeroes the group. It needs |reward| to exceed the group's own spread by about 1e14 to trip. The check lives in `combine_group` rather than its callers for the same reason the label now does: there are two of them, and a check in one is a check the other disagrees with. Raising from there also keeps the existing division of labour -- the filter, which decides whether a group trains, propagates it; the zero-std metrics turn it into "cannot tell" via `observed_reward_signal` and keep logging. The observer does not become the enforcer, and eval, which may run a different reward model entirely, is not taken down by a training-time contract. --- # ✅ Tests ## The label, not just the verdict - A flat group beginning with an unscored sample is counted *and* filed under the reward its scored samples carry. - A group with nothing scored has no label. `relax/distributed/ray/rollout.py` imports `sglang` and cannot be loaded on CPU, so these run against the shared helper -- which is why it exists. ## Both directions of the threshold - A reward with no significant digits raises, and the message names the offending component rather than an index. - A reward that is merely large (1e9 with a spread of 0.1, eight significant digits) is left alone. - The constant-sum group the floor exists for still reaches the floor and still comes out as exactly zero. - The metrics report the unreadable group as unknown while the filter refuses it. Every guard above was mutation-tested. Restoring `group[0]` reds the label tests with the production `TypeError`, not an assertion mismatch; removing the noise check reds three; tightening the threshold to 1e-18 reds ten *pre-existing* tests. The constant is bounded from both sides, not chosen. ## Verified on GPU Three Modal H100 smokes on this tree, 8 optimizer steps each: GDPO on 1 and 2 GPUs at `--global-batch-size 16` -- the first runs to reach `_whiten_by_segment`'s multi-segment path and its MAX all-reduce -- and GRPO on 1 GPU for the registry refactor's default algorithm. Both label paths produced real counts (`zero_std/count_1.0`, `count_1`, `count_0`) rather than raising, and no noise check fired on GSM8K's `{0, 1}` rewards, thirteen orders below the threshold. --- relax/agentic/rollout.py | 12 ++- relax/algorithms/rewards.py | 83 ++++++++++++++++++- relax/distributed/ray/rollout.py | 12 ++- tests/algorithms/test_gdpo.py | 74 +++++++++++++++++ .../algorithms/test_multi_reward_consumers.py | 42 +++++++++- 5 files changed, 210 insertions(+), 13 deletions(-) diff --git a/relax/agentic/rollout.py b/relax/agentic/rollout.py index 480ee0d7e..e0cdad0c8 100644 --- a/relax/agentic/rollout.py +++ b/relax/agentic/rollout.py @@ -20,7 +20,7 @@ ) from relax.agentic.profile import TRACE_KEY from relax.algorithms import get_algorithm -from relax.algorithms.rewards import metrics_group_verdict +from relax.algorithms.rewards import metrics_group_verdict, zero_std_group_label from relax.engine.filters.base_types import MetricGatherer, call_dynamic_filter from relax.engine.rollout import on_policy_distillation as opd from relax.engine.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput @@ -1311,8 +1311,14 @@ def _compute_zero_std_metrics(args, all_samples: list[Sample]) -> dict[str, floa # The distributed copy counts them instead; see `metrics_group_verdict`. if metrics_group_verdict(args, group) is not True: continue - rewarded = [sample for sample in group if sample.reward is not None] - interesting_rewards.append(str(round(rewarded[0].get_reward_value(args), 1))) + # A `True` verdict already implies a scored sample, so the label is + # never None here. It goes through the shared helper anyway: this line + # and its opposite number in the distributed copy are the half of this + # metric `metrics_group_verdict` did not cover, and they had already + # drifted -- that copy read the label off `group[0]`, scored or not. + label = zero_std_group_label(args, group) + if label is not None: + interesting_rewards.append(label) return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} diff --git a/relax/algorithms/rewards.py b/relax/algorithms/rewards.py index 206787a57..2703ac0fb 100644 --- a/relax/algorithms/rewards.py +++ b/relax/algorithms/rewards.py @@ -38,6 +38,21 @@ magnitude below the signal, so any value in this neighbourhood suppresses exactly the same things.""" +_MAX_COMPONENT_NOISE = 0.01 +"""When :func:`component_noise_scale` stops being negligible and starts being +the answer. + +The standardised values it bounds have unit variance, so this reads directly as +a fraction: at 0.01 a hundredth of what the component contributes to the +combined advantage is rounding rather than reward. That is a broken reward +scale, not a small inaccuracy -- the direction of the gradient is partly noise +whether or not the noise floor happens to zero the group. + +The fifteen orders of margin on well-conditioned rewards mean this is nowhere +near anything real: it needs |reward| to exceed the group's own spread by about +1e14 before it trips, which is a reward carrying no significant digits where it +varies.""" + def group_positions(samples: list[Any], expected_size: int) -> dict[int, list[int]]: """Map ``Sample.group_index`` to the positions it occupies in ``samples``. @@ -301,7 +316,7 @@ def component_noise_scale(group: torch.Tensor) -> torch.Tensor: return torch.where(collapsed_columns(group, dim=0), torch.zeros_like(noise), noise) -def combine_group(group: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: +def combine_group(group: torch.Tensor, weights: torch.Tensor, keys: list[str] | None = None) -> torch.Tensor: """GDPO steps 1 and 2 for one ``[G, K]`` prompt group, with a noise floor. Two components summing to a constant should standardise to exact opposites @@ -330,10 +345,41 @@ def combine_group(group: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: call; given the fifteen orders of margin on well-conditioned input, its exact value changes nothing that is not already noise. """ + noise = component_noise_scale(group) + + # Raise rather than floor when the estimate stops being negligible. The + # floor's contract is "round a residue to the zero it mathematically is", + # and that holds only while the residue *is* residue. Past + # `_MAX_COMPONENT_NOISE` the column has no significant digits left where it + # varies, and zeroing it would do precisely what + # `extract_reward_components` refuses to do one stage earlier: make a + # component that could not be read indistinguishable from one that + # genuinely collapsed, and hide a broken reward scale behind training that + # looks plausible. The failure this catches is silent in both directions -- + # groups below the floor are dropped with no log at all, and a batch where + # every group is dropped reports "rewards do not vary", which is not what + # happened. + # + # Here rather than in the callers because there are two of them -- the + # normaliser and the filter -- and a check in one of them is a check the + # other disagrees with. + if float(noise.amax()) > _MAX_COMPONENT_NOISE: + column = int(noise.argmax()) + name = f"{keys[column]!r}" if keys else f"at index {column}" + magnitude = float(group[:, column].abs().amax()) + spread = float(group[:, column].std()) + raise ValueError( + f"Reward component {name} varies by {spread:.6g} around values as large as " + f"{magnitude:.6g}, so {float(noise[column]):.2%} of its standardised value is " + f"floating-point rounding rather than reward. Rescale it (subtract the offset, or " + f"report the difference directly): standardising it divides the spread by itself, " + f"which amplifies that rounding into the advantage." + ) + standardized = standardize_group_components(group) combined = (standardized.double() * weights.double()).sum(dim=1) - floor = _NOISE_FLOOR_SAFETY * float((weights.double().abs() * component_noise_scale(group)).sum()) + floor = _NOISE_FLOOR_SAFETY * float((weights.double().abs() * noise).sum()) if float(combined.abs().amax()) <= floor: return torch.zeros_like(combined) return combined @@ -398,7 +444,7 @@ def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[fl combined = torch.zeros(len(samples), dtype=torch.float64) silent_groups = 0 for positions in positions_by_group.values(): - combined[positions] = combine_group(components[positions], weight_tensor) + combined[positions] = combine_group(components[positions], weight_tensor, keys) if not bool(combined[positions].any()): silent_groups += 1 @@ -489,7 +535,7 @@ def group_carries_reward_signal(args: Any, samples: list[Any]) -> bool: # `combine_group`, not steps 1 and 2 open-coded: this has to answer the # same question `normalize_gdpo_decoupled` will, including the noise floor. # Open-coding it is how the two came apart before. - return bool(combine_group(components, weights).any()) + return bool(combine_group(components, weights, keys).any()) def observed_reward_signal(args: Any, samples: list[Any]) -> bool | None: @@ -562,6 +608,35 @@ def metrics_group_verdict(args: Any, samples: list[Any]) -> bool | None: return None if signal is None else not signal +def zero_std_group_label(args: Any, samples: list[Any]) -> str | None: + """The ``zero_std/count_