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..42a6e6431 --- /dev/null +++ b/relax/algorithms/__init__.py @@ -0,0 +1,27 @@ +# 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, + algorithm_needs_critic, + get_algorithm, + list_algorithm_names, +) + + +__all__ = [ + "ALGORITHM_SPECS", + "AlgorithmSpec", + "algorithm_needs_critic", + "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..6b6a931f8 --- /dev/null +++ b/relax/algorithms/spec.py @@ -0,0 +1,284 @@ +# 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 + + +# The two enum-like fields below are consumed by equality checks in +# `relax/backends/megatron/loss.py` -- `advantage_normalization == "token_global"` +# at 659 and 819, `kl_level == "sequence"` at 919. Anything that is not the +# awaited string takes the *other* branch, so a typo in a registry entry does +# not fail, it silently selects a different formula. These sets are what +# `__post_init__` checks against. +KL_LEVELS = frozenset({"token", "sequence"}) +ADVANTAGE_NORMALIZATIONS = frozenset({"whiten", "token_global"}) + + +@dataclass(frozen=True) +class AlgorithmSpec: + """Everything the training pipeline needs to know about one algorithm.""" + + 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. + """ + + def __post_init__(self) -> None: + """Reject an unsupported enum value while the registry is being built. + + ``ALGORITHM_SPECS`` is a module-level literal, so this runs at import: + a typo cannot reach a worker, let alone a training step. The + implementation identifiers are already resolved eagerly for the same + reason (``_assert_spec_implementations_resolve`` in ``arguments.py``); + these two fields were the ones left unchecked, and they are the ones + whose failure is silent rather than loud -- a bad ``advantage_fn`` + raises a KeyError, a bad ``kl_level`` just trains with token-level KL + and never says so. + """ + for field, value, allowed in ( + ("kl_level", self.kl_level, KL_LEVELS), + ("advantage_normalization", self.advantage_normalization, ADVANTAGE_NORMALIZATIONS), + ): + if value not in allowed: + raise ValueError( + f"AlgorithmSpec({self.name!r}) has {field}={value!r}, which no call site matches; " + f"the run would silently take the default branch instead. " + f"Expected one of {sorted(allowed)}." + ) + + @property + def is_group_normalized(self) -> bool: + """Whether rewards get normalised per prompt group on the rollout + 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 algorithm_needs_critic(config) -> bool: + """Whether the configured algorithm runs a critic, read from the registry. + + The training pipeline asks this in five places -- role topology, the + critic's placement group, the device hand-off of the critic's ``values``, + the critic consumer's rollout fields, and the critic's own wait loop -- and + each of them used to compare ``advantage_estimator`` against ``"ppo"``. + That worked while PPO was the only value-based estimator, and silently + stopped working the moment the registry could accept a second one: + ``--advantage-estimator`` and ``ALGOS`` would take it, then none of the + value plumbing would switch on. + + ``args.use_critic`` carries the same answer, but only after + ``validate_algorithm_args`` has run; ``process_role`` and the controller's + placement logic read a config that may not have been through it yet. This + reads the spec directly so the answer does not depend on call order. + + Unknown or missing estimators answer False rather than raising: SFT and the + debug-only role paths reach these call sites with no estimator at all, and + an unknown name is rejected by argument parsing long before this matters. + """ + spec = ALGORITHM_SPECS.get(getattr(config, "advantage_estimator", None)) + return spec is not None and spec.needs_critic + + +def list_algorithm_names() -> list[str]: + """All registered algorithm names, in definition order.""" + return list(ALGORITHM_SPECS) diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..a750e4507 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -26,6 +26,7 @@ from tensordict import TensorDict from transformers import AutoConfig, AutoTokenizer +from relax.algorithms import algorithm_needs_critic from relax.distributed.checkpoint_service.client.engine import create_client from relax.distributed.ray.train_actor import TrainRayActor from relax.engine.sft.eval.runner import run_sft_eval @@ -787,7 +788,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # and land on CPU (critic ``.cpu()`` s ``values`` before PUT). Inline # GAE + normalize_advantages need GPU tensors — dispatch here so the # rest of the pipeline can assume same-device inputs. - if self.args.advantage_estimator == "ppo": + if algorithm_needs_critic(self.args): cur_device = torch.cuda.current_device() for key in ("values", "loss_masks"): tensors = rollout_data.get(key) @@ -820,7 +821,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # advantages/returns via TransferQueue; every other path (including # PPO colocate) computes GAE inline from critic's ``values``. should_compute_gae_in_actor = self.args.compute_advantages_and_returns and not ( - self.args.advantage_estimator == "ppo" and self.args.fully_async and not self.args.hybrid + algorithm_needs_critic(self.args) and self.args.fully_async and not self.args.hybrid ) if should_compute_old_log_probs: @@ -2246,7 +2247,7 @@ def _put_data_to_transfer_queue(self, output_dict=None, batch_meta=None, rollout run(self.data_system_client.async_put(data=output_dict, metadata=batch_meta)) def _put_critic_values_to_transfer_queue(self, rollout_data: RolloutBatch) -> None: - if getattr(self.args, "advantage_estimator", None) != "ppo": + if not algorithm_needs_critic(self.args): return values = rollout_data.get("values") diff --git a/relax/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/components/critic.py b/relax/components/critic.py index ef01bbf5f..14b4883f7 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -120,23 +120,25 @@ def _wait_for_rollout_data(self) -> None: self._rollout_barrier.wait_offloaded_sync() def train(self) -> None: - is_ppo = getattr(self.config, "advantage_estimator", None) == "ppo" + from relax.algorithms import algorithm_needs_critic + + has_critic = algorithm_needs_critic(self.config) while self.step < self.config.num_rollout: - if is_ppo: + if has_critic: self._wait_for_rollout_data() # In PPO colocate the actor waits for ``self.step`` to advance # past the current round before waking up, so block on training # completion here. Non-PPO critic is not on any live service graph # and keeps the historical fire-and-forget. train_ref = self.critic_model.async_train(self.step) - if is_ppo: + if has_critic: ray.get(train_ref) # Note: save_model runs inside ``train_critic`` (backend) while the # model is still awake, so no explicit save call here. # In critic-only warmup, actor+advantages never consume the partition, # so critic must clear it itself; steady-state clearing stays with actor. - if is_ppo and self.step < getattr(self.config, "num_critic_only_steps", 0): + if has_critic and self.step < getattr(self.config, "num_critic_only_steps", 0): run(self.data_system_client.async_clear_partition(partition_id=f"train_{self.step}")) try: diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..bf78dda14 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -30,6 +30,7 @@ deploy_agentic_chat_api_services, shutdown_agentic_chat_api_services, ) +from relax.algorithms import algorithm_needs_critic from relax.core.optional_roles import register_extra_roles from relax.core.registry import ALGOS, ROLES, process_role from relax.core.service import Service, create_placement_group @@ -104,7 +105,7 @@ def _actor_rollout_pg_roles(config: Namespace) -> list[str]: matches the actor's; otherwise critic runs on its own placement group. """ roles = list(ACTOR_ROLLOUT_PG_ROLES) - if getattr(config, "advantage_estimator", None) != "ppo": + if not algorithm_needs_critic(config): return roles resource = getattr(config, "resource", None) or {} if resource.get("critic") == resource.get("actor"): diff --git a/relax/core/registry.py b/relax/core/registry.py index a295e3bcb..e70c0316e 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, algorithm_needs_critic from relax.components.actor import Actor from relax.components.actor_fwd import ActorFwd from relax.components.advantages import Advantages @@ -80,68 +81,41 @@ 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, which reads the same + ``needs_critic`` through ``algorithm_needs_critic``. ``controller.py`` + iterates ``list(process_role(config))`` and + skips any role missing from this dict, so an algorithm without a critic + simply never matches the ``critic`` member the role sets already carry. + """ + 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, } @@ -152,7 +126,7 @@ def process_role(config): return ROLES_TRAIN_ONLY if getattr(config, "loss_type", None) == "sft": return ROLES_SFT_ONLY - if getattr(config, "advantage_estimator", None) == "ppo": + if algorithm_needs_critic(config): if config.fully_async: if getattr(config, "true_on_policy_mode", False): return ROLES_PPO_FULLY_ASYNC_ON_POLICY diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index f7d1e710e..c7313e54b 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,280 @@ 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. + # + # A YAML file that *names* `global_batch_size` is a different case: that is + # not a stale value left over from an earlier derivation, it is the override + # this function exists to apply. Re-deriving over it wrote the YAML's value + # and then replaced it in the next statement, so the run used neither the + # configured number nor an error -- the one outcome the override contract + # rules out. Derive first so the comparison has something to name, then + # refuse the conflict rather than picking a winner. + yaml_global_batch_size = data.get("global_batch_size") + derive_global_batch_size(args, enforce_consistency=False) + if yaml_global_batch_size is not None and args.global_batch_size != yaml_global_batch_size: + raise ValueError( + f"--custom-config-path sets global_batch_size to {yaml_global_batch_size}, but " + f"num_steps_per_rollout {args.num_steps_per_rollout} over rollout_batch_size " + f"{args.rollout_batch_size} * n_samples_per_prompt {args.n_samples_per_prompt} derives " + f"{args.global_batch_size}. Remove one of the two from the YAML -- whichever you drop, " + f"the other is what the run would otherwise have used without saying so." + ) + validate_batch_shape(args) + if args.use_critic != use_critic_before_override: + # Role composition and the offload flags were derived from the pre-override + # 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 +3324,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 +3409,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 +3426,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 +3578,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 +3747,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 +3827,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/training/data_fields.py b/relax/utils/training/data_fields.py index ebd6479af..b7aa888d3 100644 --- a/relax/utils/training/data_fields.py +++ b/relax/utils/training/data_fields.py @@ -32,12 +32,14 @@ def build_data_fields(args: Namespace, *, consumer: str = "actor") -> list[str]: fields.append("multimodal_train_inputs") return fields - is_ppo = getattr(args, "advantage_estimator", None) == "ppo" + from relax.algorithms import algorithm_needs_critic - if is_ppo and consumer == "critic": + has_critic = algorithm_needs_critic(args) + + if has_critic and consumer == "critic": return _base_rollout_fields(args) - if is_ppo and consumer == "advantages": + if has_critic and consumer == "advantages": # PPO colocate never runs actor_fwd, so ref/log_probs are only # requested when kl_coef != 0 (i.e. an actor_fwd role is present). fields = _base_rollout_fields(args) @@ -53,7 +55,7 @@ def build_data_fields(args: Namespace, *, consumer: str = "actor") -> list[str]: return fields fields = _base_rollout_fields(args) - if is_ppo: + if has_critic: # PPO colocate: actor consumes critic's ``values`` and computes GAE # inline. Fully_async: standalone Advantages service produces # ``advantages``/``returns``, actor just pulls the finished tensors. diff --git a/relax/utils/training/ppo_utils.py b/relax/utils/training/ppo_utils.py index e4eab1310..b5bd4790f 100644 --- a/relax/utils/training/ppo_utils.py +++ b/relax/utils/training/ppo_utils.py @@ -11,6 +11,7 @@ import torch.distributed as dist import torch.nn.functional as F +from relax.algorithms import algorithm_needs_critic from relax.utils.logging_utils import get_logger @@ -18,12 +19,13 @@ def validate_ppo_config(config: Namespace) -> None: - if getattr(config, "advantage_estimator", None) != "ppo": + if not algorithm_needs_critic(config): return resource = getattr(config, "resource", None) or {} if "critic" not in resource: - raise ValueError("--advantage-estimator ppo requires a 'critic' entry in --resource.") + estimator = getattr(config, "advantage_estimator", None) + raise ValueError(f"--advantage-estimator {estimator} requires a 'critic' entry in --resource.") if getattr(config, "fully_async", False) or getattr(config, "hybrid", False): raise ValueError("PPO does not currently support --fully-async or --hybrid.") diff --git a/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..116c0d19a --- /dev/null +++ b/tests/algorithms/test_algorithm_registry.py @@ -0,0 +1,200 @@ +# 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 + + +# ---------------- enum-like fields must fail loudly, not silently ---------------- + + +@pytest.mark.parametrize( + "field,bad", + [ + ("kl_level", "Sequence"), # right word, wrong case + ("kl_level", "seq"), + ("advantage_normalization", "token-global"), # hyphen instead of underscore + ("advantage_normalization", "none"), + ], +) +def test_an_unsupported_enum_value_is_refused_when_the_spec_is_built(field, bad): + """These two fields are compared for equality, so a typo picks a formula. + + `loss.py` asks for `advantage_normalization == "token_global"`, and for + `kl_level == "sequence"`; every other string takes the else branch. Unlike + `advantage_fn`, which blows up with a KeyError the first time it is looked + up, a misspelled value here starts training successfully and quietly uses + the wrong KL level or the wrong advantage normalisation. The registry is a + module-level literal, so validating in `__post_init__` moves that from a + silent wrong-maths run to an import-time error. + """ + from relax.algorithms.spec import AlgorithmSpec + + kwargs = dict( + name="probe", + reward_normalizer="none", + advantage_fn="grpo_broadcast", + policy_loss_fn="ppo_clip", + ) + kwargs[field] = bad + + with pytest.raises(ValueError, match=field): + AlgorithmSpec(**kwargs) + + +def test_the_supported_values_are_the_ones_the_shipped_specs_use(): + """The allow-list must not drift from the registry it guards.""" + from relax.algorithms.spec import ADVANTAGE_NORMALIZATIONS, ALGORITHM_SPECS, KL_LEVELS + + assert {s.kl_level for s in ALGORITHM_SPECS.values()} <= KL_LEVELS + assert {s.advantage_normalization for s in ALGORITHM_SPECS.values()} <= ADVANTAGE_NORMALIZATIONS diff --git a/tests/algorithms/test_algos_roles.py b/tests/algorithms/test_algos_roles.py new file mode 100644 index 000000000..47e2e5b7c --- /dev/null +++ b/tests/algorithms/test_algos_roles.py @@ -0,0 +1,230 @@ +# 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" + + +# ---------------- a second critic algorithm must need no new name checks ---------------- + + +def _register_second_critic_algorithm(monkeypatch): + """Add a value-based estimator to the registry for the duration of a test. + + Returns its name. Uses PPO's own spec so the only thing under test is + whether the pipeline keys off `needs_critic` or off the literal `"ppo"`. + """ + import dataclasses + + from relax.algorithms import ALGORITHM_SPECS + from relax.algorithms import spec as spec_module + + name = "ppo_second" + clone = dataclasses.replace(ALGORITHM_SPECS["ppo"], name=name) + monkeypatch.setitem(spec_module.ALGORITHM_SPECS, name, clone) + assert ALGORITHM_SPECS[name].needs_critic is True + return name + + +def _args_for(name, **overrides): + from argparse import Namespace + + base = dict( + advantage_estimator=name, + multimodal_keys=None, + kl_coef=0.0, + fully_async=False, + hybrid=False, + use_opd=False, + use_rollout_logprobs=False, + true_on_policy_mode=False, + debug_rollout_only=False, + debug_train_only=False, + loss_type=None, + ) + base.update(overrides) + return Namespace(**base) + + +def test_a_second_critic_algorithm_gets_the_critic_rollout_fields(monkeypatch): + """`values` must reach the advantages consumer, unprompted by any name. + + This is the failure the registry was supposed to make impossible: argparse + and `ALGOS` accept a second value-based estimator, and then the value + plumbing silently does not switch on because it compares against `"ppo"`. + """ + from relax.utils.training.data_fields import build_data_fields + + name = _register_second_critic_algorithm(monkeypatch) + + ppo_fields = build_data_fields(_args_for("ppo"), consumer="advantages") + new_fields = build_data_fields(_args_for(name), consumer="advantages") + assert new_fields == ppo_fields, "a second critic algorithm sees different fields than PPO" + assert "values" in new_fields + + # and the critic consumer's own set, which is the base set rather than the actor's + assert build_data_fields(_args_for(name), consumer="critic") == build_data_fields( + _args_for("ppo"), consumer="critic" + ) + + +def test_a_second_critic_algorithm_is_told_it_needs_a_critic_resource(monkeypatch): + """Startup validation is keyed on the capability, not on the name.""" + import pytest as _pytest + + from relax.utils.training.ppo_utils import validate_ppo_config + + name = _register_second_critic_algorithm(monkeypatch) + + with _pytest.raises(ValueError, match="requires a 'critic' entry"): + validate_ppo_config(_args_for(name, resource={"actor": "a"})) + + # a non-critic estimator is still waved through + validate_ppo_config(_args_for("grpo", resource={"actor": "a"})) + + +@requires_megatron +def test_a_second_critic_algorithm_walks_the_critic_role_topology(monkeypatch): + """`process_role` decides which roles the controller walks at all.""" + from relax.core.registry import process_role + + name = _register_second_critic_algorithm(monkeypatch) + + # identity, not membership: every role set carries a `critic` member and + # `ALGOS` is what filters it out, so comparing member names would pass even + # if the second algorithm fell through to the non-critic topology. + assert process_role(_args_for(name)) is process_role(_args_for("ppo")) + assert process_role(_args_for(name)) is not process_role(_args_for("grpo")) + + # and the fully-async split follows too, rather than only the colocate one + assert process_role(_args_for(name, fully_async=True)) is process_role(_args_for("ppo", fully_async=True)) diff --git a/tests/algorithms/test_arguments_spec_driven.py b/tests/algorithms/test_arguments_spec_driven.py new file mode 100644 index 000000000..ab6798eb4 --- /dev/null +++ b/tests/algorithms/test_arguments_spec_driven.py @@ -0,0 +1,372 @@ +# 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) + + +# ---------------- a YAML global_batch_size must not be derived over ---------------- + + +def _batch_args(tmp_path, body, **overrides): + """Args shaped for the batch-size derivation, already validated once.""" + base = _overridable_args(tmp_path, body, **overrides) + base.rollout_batch_size = 32 + base.n_samples_per_prompt = 4 + base.num_steps_per_rollout = 4 + base.global_batch_size = 32 # 32 * 4 // 4, i.e. what the pre-merge derivation wrote + base.micro_batch_size = 1 + base.use_dynamic_batch_size = False + for key, value in overrides.items(): + setattr(base, key, value) + return base + + +def test_yaml_global_batch_size_conflicting_with_the_derivation_is_refused(arguments_module, tmp_path): + """The override used to be written and then silently replaced. + + `apply_custom_config_overrides` merges the YAML, then re-derives + `global_batch_size` from `num_steps_per_rollout`. With the derivation + unconditional, a YAML that names `global_batch_size` had its value assigned + by the merge loop and overwritten one statement later -- so the run used + neither the configured number nor an error, which is the single outcome the + "YAML key overrides the argument" contract does not allow. + """ + args = _batch_args(tmp_path, "global_batch_size: 999\n") + + with pytest.raises(ValueError, match="sets global_batch_size to 999"): + arguments_module.apply_custom_config_overrides(args) + + +def test_yaml_global_batch_size_agreeing_with_the_derivation_survives(arguments_module, tmp_path): + """Naming the value the derivation would reach anyway is not a conflict.""" + args = _batch_args(tmp_path, "global_batch_size: 32\n") + + arguments_module.apply_custom_config_overrides(args) + + assert args.global_batch_size == 32 + + +def test_yaml_that_only_moves_a_derivation_input_still_re_derives(arguments_module, tmp_path): + """The behaviour the `enforce_consistency=False` call was added for. + + Switching `num_steps_per_rollout` from 4 to 1 has to produce + `rollout * n`; the pre-merge value of 32 is stale by construction and must + not be compared against. + """ + args = _batch_args(tmp_path, "num_steps_per_rollout: 1\n") + + arguments_module.apply_custom_config_overrides(args) + + assert args.global_batch_size == 128 # 32 * 4 // 1 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..9b2e77c56 --- /dev/null +++ b/tests/algorithms/test_dispatch_parity_vs_main.py @@ -0,0 +1,643 @@ +# 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) + + +# ---------------- GAE: the adapter main's two call sites disagreed on ---------------- +# +# `advantage_gae` had only a co_names check, and it is the adapter with the most +# to get wrong: it shapes the reward in place before delegating, and it carries +# `padded_total_lengths`, which one of main's two call sites passed and the +# other did not. A dropped or reordered argument here does not raise -- it reads +# the wrong token positions and trains on them. + + +def _gae_inputs(): + """Fresh tensors per call: `advantage_gae` mutates `kl` in place.""" + kl = [torch.tensor([0.1, 0.2, 0.3]), torch.tensor([0.4, 0.5])] + values = [torch.tensor([0.5, 0.25, 0.125]), torch.tensor([1.0, 2.0])] + return dict( + rewards=[1.5, -2.0], + kl=kl, + values=values, + response_lengths=[3, 2], + total_lengths=[3, 2], + ) + + +def _main_gae(kl_coef, gamma, lambd, *, padded_total_lengths=None): + """main's PPO branch, transcribed, not regenerated. + + components/advantages.py:181-193 and the megatron duplicate at + loss.py:585-602. The only difference between the two is that the megatron + one forwards `padded_total_lengths`; both shape the reward identically. + """ + inputs = _gae_inputs() + old_rewards, kl = inputs["rewards"], inputs["kl"] + rewards = [] + for reward, k in zip(old_rewards, kl, strict=False): + k *= -kl_coef + cp_rank = 0 # the fixture pins mpu.get_context_parallel_rank() to 0 + if cp_rank == 0: + k[-1] += reward + rewards.append(k) + return ppo_utils.get_advantages_and_returns_batch( + inputs["total_lengths"], + inputs["response_lengths"], + inputs["values"], + rewards, + gamma, + lambd, + padded_total_lengths=padded_total_lengths, + ) + + +@pytest.mark.parametrize( + "kl_coef,gamma,lambd", + [ + (0.0, 1.0, 1.0), # the degenerate case the co_names check implied was enough + (0.05, 0.99, 0.95), # non-zero kl_coef and real discounting + ], +) +def test_gae_adapter_matches_mains_numbers(cp_disabled, kl_coef, gamma, lambd): + """Every element, not just the kernel's name in the bytecode.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + args = _args("ppo", kl_coef=kl_coef, gamma=gamma, lambd=lambd) + got_adv, got_ret = compute_advantages_and_returns(args, **_gae_inputs()) + want_adv, want_ret = _main_gae(kl_coef, gamma, lambd) + + for got, want, label in ((got_adv, want_adv, "advantages"), (got_ret, want_ret, "returns")): + assert len(got) == len(want), label + for left, right in zip(got, want, strict=True): + torch.testing.assert_close(left, right, rtol=0, atol=0, msg=label) + + +def test_gae_terminal_reward_lands_on_the_last_token(cp_disabled): + """The terminal reward is the whole signal; without it, only KL trains.""" + from relax.algorithms.advantages import compute_advantages_and_returns + + args = _args("ppo", kl_coef=0.0, gamma=1.0, lambd=1.0) + with_reward, _ = compute_advantages_and_returns(args, **_gae_inputs()) + + zeroed = _gae_inputs() + zeroed["rewards"] = [0.0, 0.0] + without_reward, _ = compute_advantages_and_returns(args, **zeroed) + + assert not torch.equal(with_reward[0], without_reward[0]), "the terminal reward changed nothing" + + +def test_gae_adapter_forwards_padded_total_lengths_to_the_kernel(cp_disabled, monkeypatch): + """The argument the two call sites disagreed on must survive the adapter. + + Scope, stated because it is easy to over-read: `padded_total_lengths` is + only *consumed* when `cp_size > 1` (ppo_utils.py, the `all_gather_with_cp` + branch), and this file runs at cp_size 1. So this pins that the adapter + hands the value through unchanged and in the right keyword -- the "dropped + or reordered argument" failure -- not that the padded slicing itself is + correct. That needs a real context-parallel group and is not covered here. + """ + from relax.algorithms import advantages as advantages_module + + seen = {} + + def spy(*args, **kwargs): + seen["args"] = args + seen["kwargs"] = kwargs + return ([torch.zeros(3)], [torch.zeros(3)]) + + monkeypatch.setattr(advantages_module, "get_advantages_and_returns_batch", spy) + + padded = [8, 8] + advantages_module.compute_advantages_and_returns( + _args("ppo", kl_coef=0.0, gamma=1.0, lambd=1.0), + **_gae_inputs(), + padded_total_lengths=padded, + ) + + assert seen["kwargs"].get("padded_total_lengths") == padded + # main's positional order: total_lengths, response_lengths, values, rewards, gamma, lambd + assert seen["args"][0] == [3, 2] + assert seen["args"][1] == [3, 2] + assert seen["args"][4] == 1.0 and seen["args"][5] == 1.0 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