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/examples/algorithms.md b/docs/en/examples/algorithms.md index 370e5adda..93ccaf48d 100644 --- a/docs/en/examples/algorithms.md +++ b/docs/en/examples/algorithms.md @@ -277,6 +277,91 @@ SAPO_ARGS=( --- +## GDPO + +GDPO (Group reward-Decoupled Normalization Policy Optimization, [arXiv 2601.05242](https://arxiv.org/abs/2601.05242)) targets **multi-reward** training. It standardizes each reward component within its prompt group and only then combines them, instead of summing the rewards first and normalizing once as GRPO does. + +### How It Works + +For prompt $i$ with $G$ rollouts and $n$ reward components: + +**Step 1 — per-reward group standardization:** + +$$A_k^{(i,j)} = \frac{r_k^{(i,j)} - \mathrm{mean}_j\{r_k^{(i,\cdot)}\}}{\mathrm{std}_j\{r_k^{(i,\cdot)}\} + \epsilon}$$ + +**Step 2 — weighted sum:** + +$$A_\text{sum}^{(i,j)} = \sum_k w_k A_k^{(i,j)}$$ + +The weights multiply the **normalized** advantages, not the raw rewards. After step 1 every component is on the same scale, so a weight expresses relative importance rather than the component's units. + +**Step 3 — batch-wise whitening:** + +$$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\mathrm{std}_\text{batch} + \epsilon}$$ + +**Why this beats GRPO:** summing first and standardizing once discards two things. + +*Correlation structure between components.* GRPO standardizes the sum, so every group comes out at unit variance whether its components corroborate or contradict each other. GDPO gives each component unit variance first, so the combined variance is `Σwᵢ² + 2Σwᵢwⱼρᵢⱼ` — `2 + 2ρ` for two equal weights — and is therefore **decided by the correlation**: `ρ→+1` amplifies (measured 2x), `ρ=0` gives √2, `ρ→−1` attenuates to zero. Step 3 whitens across the *batch*, so that between-group difference reaches the final advantage. Measured (G=2, ρ=+1, two groups in one batch): GRPO gives both ±0.707; GDPO gives ±0.548 and ±1.095. + +This is **not** "more varying components means more signal" — that holds only for `ρ>0`. At `ρ=−0.8` the combined signal is 0.63x a single component. + +*Scale disparity between components.* A `correctness` in {0, 1} added to a reward in the hundreds (the paper's maths setup scores response length) yields a sum whose variance is essentially the large component's, so GRPO's direction is decided by it alone. GDPO gives each component unit variance first, so a weight expresses relative importance rather than units. + +**What GDPO does not do:** rescue a group whose components sum to a constant. There `r₂ = C − r₁` forces `z₂ = −z₁`, so equal weights cancel to exactly zero — the same answer GRPO gives. Only unequal weights break that tie. If *every* component is constant, GDPO returns zero as well. + +**On $\epsilon$:** GDPO uses $\epsilon = 10^{-4}$ at both steps, matching the reference implementation (the `scale_rewards` GDPO branch of TRL's `GRPOTrainer`), whereas GRPO / GSPO / SAPO / CISPO keep this repository's existing $10^{-6}$. The two only diverge on near-degenerate groups: with binary rewards and a group of 8 the within-group standard deviation is around 0.4 and the constants differ by 0.02%, but a continuous reward (the paper's maths setup scores response length) can leave a group at a standard deviation of ~$10^{-3}$, where $10^{-4}$ damps that group's signal by about 7% against 0.08% for $10^{-6}$. Groups that collapse *exactly* never reach this division; they are detected by exact equality and zeroed. + +### Key Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--advantage-estimator gdpo` | — | Enable GDPO | +| `--gdpo-reward-keys` | — | **Required**, at least two. Keys in the reward dict to standardize independently, e.g. `correctness format` | +| `--gdpo-reward-weights` | all 1.0 | Per-component weights; length must match `--gdpo-reward-keys` | +| `--reward-key` | — | **Required**; selects the scalar used for metrics and the `raw_reward` column | +| `--n-samples-per-prompt` | — | Must be >= 2 (the unbiased group std is undefined at $G=1$) | + +The reward function must return a dict containing every key. A missing key, a non-numeric value, a bool, or NaN/Inf raises rather than defaulting to 0.0 — a silently zeroed component is indistinguishable from a genuinely collapsed one. + +### Quick Start + +```bash +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --gdpo-reward-weights 1.0 1.0 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + --n-samples-per-prompt 8 +) +``` + +A complete runnable example lives in [`examples/gdpo/`](https://github.com/redai-infra/Relax/tree/main/examples/gdpo). + +### Known Deviations + +Two differences between this implementation and the paper. Confirm they are acceptable before training. Step 3's batch boundary used to be a third; it has since been corrected — see below. + +**Step 3's batch boundary (now aligned).** Eq. 6 normalises over one training batch. The caller merges `num_rollout_minis` of them with `concat_rollout_batches` before the advantage stage, so step 3 has to be told where the boundaries are. They travel in `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY`, which all three actor paths (colocate and hybrid) set; `loss.py` forwards them to the advantage dispatcher as `mini_batch_sizes`, and **only GDPO reads it** — every other estimator absorbs it in `**_unused` and is bit-identical either way. Each segment all-reduces across the data-parallel group, so the statistics cover both a whole training batch and every rank. Why it matters: whitening merged batches centres them all on a pooled mean, and on a measured example four of eight samples **change sign** — a different objective, not a precision difference. That is also why absent boundaries are an **error** rather than a fallback to merged whitening: a caller that omits the metadata would optimise the wrong objective with loss and grad_norm both fine. + +**`--fully-async` remains unsupported** and is rejected during argument validation: it hands advantage computation to the single-replica Advantages deployment, which has no data-parallel group, never sees the batch boundaries, and consumes one `global_batch_size / num_iters_per_train_update` slice at a time — when that quotient is 1 the whitened output is identically zero and the run trains on no signal at all, quietly. + +1. **A single reward does not reduce to GRPO.** Step 3 still applies, leaving a positive scalar difference from GRPO (data-dependent, measured around 1.21). Use `--advantage-estimator grpo` if you want GRPO semantics. +2. **$G=2$ discards magnitude.** Any two distinct values standardize to exactly $\pm 1/\sqrt{2}$, so with a group of two the only thing distinguishing components is their weights. + +### Mutually Exclusive Options + +- `--normalize-advantages`: step 3 already whitens per sequence; adding the token-level pass on top is not meaningful. +- `--custom-reward-post-process-path`: that hook short-circuits reward post-processing entirely, silently skipping steps 1 and 2 while the run still reports itself as GDPO. +- `--agentic-custom-advantage-path`: the second early return in `post_process_rewards`, which likewise returns ahead of the normalizer, with the same consequence. One flag, `AlgorithmSpec.allows_reward_post_process_hooks`, guards both. +- `--fully-async`: see above. + +All of these fail during argument validation. + +`--dynamic-sampling-filter-path` does **not** conflict: the built-in `check_reward_nonzero_std` is component-aware, computing what GDPO's first two steps actually produce and keeping the group only when that is non-zero, so its verdict matches the signal training receives. A warning is logged only for a *custom* filter, which may reduce the group to the single `--reward-key` scalar and drop groups whose signal lives in the other components. + +--- + ## Algorithm Comparison | Algorithm | Advantage Computation | Policy Loss | KL Constraint | @@ -289,6 +374,7 @@ SAPO_ARGS=( | **GSPO** | Group-relative reward | PPO-Clip + sequence-level KL | Sequence-level ratio | | **SAPO** | Group-relative reward | Sigmoid gate | Temperature-controlled | | **RLOO** | Leave-one-out baseline | Unclipped REINFORCE | Optional KL loss (same as GRPO) | +| **GDPO** | Per-reward group standardization + weighted sum + batch whitening | PPO-Clip (hard clip) | Optional KL loss | ## Next Steps diff --git a/docs/en/guide/adding-an-algorithm.md b/docs/en/guide/adding-an-algorithm.md new file mode 100644 index 000000000..be5e8f083 --- /dev/null +++ b/docs/en/guide/adding-an-algorithm.md @@ -0,0 +1,183 @@ +# 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 +└── numerics.py shared numeric constants and degeneracy guards +``` + +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 (as GDPO needs `--gdpo-reward-keys`) | 4-6 (the above, plus the option and its validation in `arguments.py`, plus an example and docs) | + +What the registry removes is one algorithm name being interpreted in six +scattered if/elif chains — not the cost of adding an algorithm. GDPO is 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 | +| `supports_fully_async` | Set `False` to reject `--fully-async`, where advantages are computed slice-by-slice in a single-replica service with no data-parallel group | +| `allows_reward_post_process_hooks` | Set `False` to block both `--custom-reward-post-process-path` and `--agentic-custom-advantage-path`; each returns from `post_process_rewards` ahead of the normalizer and would silently skip your reward stage | +| `uses_reward_components` | The algorithm consumes several named reward components rather than one scalar; drives the `--gdpo-reward-keys` validation | + +The four `validate_*` functions in `relax/utils/arguments.py` consume every +field in that table except `kl_level`, `needs_full_log_probs` and +`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 — multi-reward algorithms such as GDPO collapse their +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) +- GDPO is the most recent algorithm to go through this process; read + `relax/algorithms/` alongside `examples/gdpo/`. diff --git a/docs/en/guide/configuration.md b/docs/en/guide/configuration.md index f98a14739..0e9d12545 100644 --- a/docs/en/guide/configuration.md +++ b/docs/en/guide/configuration.md @@ -299,7 +299,7 @@ bash scripts/training/text/run-qwen3-4B-fp16-8xgpu.sh \ | Parameter | Type | Default | Options | Description | |-----------|------|---------|---------|-------------| -| `--advantage-estimator` | str | grpo | `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`, `gdpo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline` | Advantage estimator. `--help` is authoritative: the choices are read from the registry, so a new algorithm appears there without this table being edited. OPD is independent of this choice; enable it with `--use-opd` and its KL/loss coefficient | | `--normalize-advantages` | flag | False | - | Whether to normalize advantages | | `--disable-grpo-std-normalization` | flag | - | - | Disable GRPO standard deviation normalization (from [Dr.GRPO](https://arxiv.org/pdf/2503.20783)) | | `--disable-rewards-normalization` | flag | - | - | Disable reward normalization | diff --git a/docs/zh/examples/algorithms.md b/docs/zh/examples/algorithms.md index 76a8b8231..f1e38237b 100644 --- a/docs/zh/examples/algorithms.md +++ b/docs/zh/examples/algorithms.md @@ -274,6 +274,91 @@ SAPO_ARGS=( --- +## GDPO + +GDPO（Group reward-Decoupled Normalization Policy Optimization，[arXiv 2601.05242](https://arxiv.org/abs/2601.05242)）面向**多奖励**训练。它对每个奖励分量分别做组内标准化，再合并——而不是像 GRPO 那样先把多个奖励加起来再归一化。 + +### 算法原理 + +设第 $i$ 个 prompt 采样 $G$ 条 rollout，共 $n$ 个奖励分量。 + +**第一步 —— 逐奖励组内标准化**： + +$$A_k^{(i,j)} = \frac{r_k^{(i,j)} - \mathrm{mean}_j\{r_k^{(i,\cdot)}\}}{\mathrm{std}_j\{r_k^{(i,\cdot)}\} + \epsilon}$$ + +**第二步 —— 加权求和**： + +$$A_\text{sum}^{(i,j)} = \sum_k w_k A_k^{(i,j)}$$ + +注意权重乘在**归一化后的 advantage** 上，不是乘在原始 reward 上。经过第一步各分量已在同一尺度，权重表达的是相对重要性，而不是分量的量纲。 + +**第三步 —— batch 级白化**： + +$$\hat{A}^{(i,j)} = \frac{A_\text{sum}^{(i,j)} - \mathrm{mean}_\text{batch}}{\mathrm{std}_\text{batch} + \epsilon}$$ + +**相对 GRPO 的收益**：GRPO 把各分量相加后只做一次组内标准化，这会丢掉两类信息。 + +*一是分量间的相关结构*：GRPO 对和做组内标准化，每组输出恒为单位方差，无论分量彼此印证还是互相矛盾。GDPO 让各分量先各自标准化再加权求和，组合方差为 `Σwᵢ² + 2Σwᵢwⱼρᵢⱼ`（等权重两分量时为 `2 + 2ρ`），**由相关系数决定**：`ρ→+1` 增强（实测 2 倍）、`ρ=0` 为 √2 倍、`ρ→−1` 减弱直至归零。第三步的 batch 白化跨组进行，这个组间强度差异会保留到最终 advantage。实测（G=2、ρ=+1，一个 batch 两组）：GRPO 两组同为 ±0.707，GDPO 分别为 ±0.548 与 ±1.095。 + +注意这条**不是**「变化的分量越多信号越强」——那只在 `ρ>0` 时成立。`ρ=−0.8` 时组合信号反而只有单分量的 0.63 倍。 + +*二是分量间的尺度差异*：`correctness ∈ {0,1}` 与一个取值上百的分量（论文实验用响应长度）相加时，和的方差几乎全部来自后者，GRPO 的方向由它单独决定；GDPO 先让每个分量单位方差，权重才真正表达相对重要性而非量纲。 + +**GDPO 做不到什么**：若各分量在组内**恰好加和为常数**（`r₂ = C − r₁`），标准化后恒有 `z₂ = −z₁`，等权重下**完全抵消为零**，与 GRPO 结果相同——只有不等权重能在这类组上取得信号。若**所有**分量都恒定，GDPO 同样返回零。 + +**关于 $\epsilon$**：GDPO 的两步都用 $\epsilon = 10^{-4}$，与参考实现（TRL `GRPOTrainer` 的 `scale_rewards` GDPO 分支）一致，而 GRPO / GSPO / SAPO / CISPO 沿用本仓库既有的 $10^{-6}$。两者的差别只在近乎塌缩的组上显现：二值 reward、组大小 8 时组内标准差约 0.4，两个取值的差异是 0.02%；但连续 reward（论文的数学实验用响应长度）可能让某组的标准差落到 $10^{-3}$ 量级，此时 $10^{-4}$ 会把该组的信号额外压低约 7%，而 $10^{-6}$ 只压低 0.08%。**完全**塌缩的组不会走到这个除法——它们由 exact 相等判定后直接置零。 + +### 关键参数 + +| 参数 | 默认值 | 说明 | +|------|--------|------| +| `--advantage-estimator gdpo` | — | 启用 GDPO | +| `--gdpo-reward-keys` | — | **必填**，至少两个。奖励函数返回的 dict 中要独立归一化的 key，如 `correctness format` | +| `--gdpo-reward-weights` | 全 1.0 | 各分量权重，长度须与 `--gdpo-reward-keys` 一致 | +| `--reward-key` | — | **必填**，选出用于 metrics 与 `raw_reward` 列的标量 | +| `--n-samples-per-prompt` | — | 必须 ≥ 2（组内无偏标准差在 $G=1$ 时无定义） | + +奖励函数必须返回包含全部 key 的 dict。缺 key、非数值、bool、NaN/Inf 都会直接报错而不是填 0——静默填 0 会把契约违约伪装成真实的 reward collapse。 + +### 快速开始 + +```bash +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --gdpo-reward-weights 1.0 1.0 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + --n-samples-per-prompt 8 +) +``` + +完整可运行示例见 [`examples/gdpo/`](https://github.com/redai-infra/Relax/tree/main/examples/gdpo)。 + +### 已知偏差 + +以下两点是实现与论文之间的实际差异，训练前请确认可以接受。第三步的 batch 边界曾经也在此列，现已修正——见下。 + +**第三步的 batch 边界（已对齐）**。论文 Eq. 6 在**一个训练批**上归一化。调用方会先把 `num_rollout_minis` 个训练批用 `concat_rollout_batches` 合并再进 advantage 阶段，所以第三步必须知道批边界。边界由 `ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY` 携带（colocate 与 hybrid 三条路径都写入），`loss.py` 作为 `mini_batch_sizes` 传给 advantage 分发器，**只有 GDPO 消费**——其余估计器由 `**_unused` 吞掉，逐位不变。每段各自跨 DP all-reduce，所以统计量既覆盖完整训练批、也覆盖全部 rank。这一点为什么重要：合并白化会把两个批都对着共同均值中心化，实测 8 个样本里有 4 个**符号翻转**——那是另一个优化目标，不是精度差异。正因如此，边界缺失时 GDPO **直接报错**而不是退回合并白化：一个漏写这份元数据的调用方会在 loss、grad_norm 全部正常的情况下优化错误的目标。 + +**`--fully-async` 仍不受支持**，参数校验阶段直接拒绝：那条路径把 advantage 计算交给单副本的 Advantages 服务，它没有数据并行通信域，也拿不到批边界，且每次只消费 `global_batch_size / num_iters_per_train_update` 的一个切片；当这个商为 1 时白化输出恒为 0，训练会安静地在零信号上跑完。 + +1. **单个奖励时 GDPO 不退化为 GRPO**。第三步仍然生效，结果与 GRPO 相差一个正标量（与数据相关，实测约 1.21）。要 GRPO 语义就直接用 `--advantage-estimator grpo`。 +2. **$G=2$ 时幅度信息丢失**。任意两个不同值经无偏标准化后恒为 $\pm 1/\sqrt{2}$，此时分量之间的区分度只来自权重。 + +### 互斥项 + +- 不能与 `--normalize-advantages` 同用：第三步已经做过序列级白化，再叠加 token 级白化没有意义。 +- 不能与 `--custom-reward-post-process-path` 同用：该钩子会整段短路奖励后处理，导致第一、二步被静默跳过，而训练日志仍然显示算法是 GDPO。 +- 不能与 `--agentic-custom-advantage-path` 同用：`post_process_rewards` 里的第二个早返回点，同样赶在归一化器之前返回，后果与上一条相同。两者由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守。 +- 不能与 `--fully-async` 同用（见上）。 + +以上都会在参数校验阶段直接报错。 + +`--dynamic-sampling-filter-path` **不**冲突：内置的 `check_reward_nonzero_std` 已经是分量感知的，它直接算出 GDPO 前两步的组合结果、按其是否非零判定，因此与训练实际拿到的信号一致。只有指向**自定义** filter 时才会给出警告——那种 filter 若只看 `--reward-key` 标量，就会丢掉只存在于其它分量的信号。 + +--- + ## 算法对比 | 算法 | Advantage 计算 | 策略损失 | KL 约束方式 | @@ -286,6 +371,7 @@ SAPO_ARGS=( | **GSPO** | 组相对奖励 | PPO-Clip + 序列级 KL | 序列级 ratio | | **SAPO** | 组相对奖励 | Sigmoid 门控 | 温度控制 | | **RLOO** | Leave-one-out 基线 | 非裁剪 REINFORCE | 可选 KL loss（同 GRPO） | +| **GDPO** | 逐奖励组内标准化 + 加权求和 + batch 白化 | PPO-Clip（硬裁剪） | 可选 KL loss | ## 下一步 diff --git a/docs/zh/guide/adding-an-algorithm.md b/docs/zh/guide/adding-an-algorithm.md new file mode 100644 index 000000000..88212d77f --- /dev/null +++ b/docs/zh/guide/adding-an-algorithm.md @@ -0,0 +1,132 @@ +# 接入一个新算法 + +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 +└── numerics.py 共享的数值常量与退化判定 +``` + +三条硬约束： + +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` + 对应的实现模块） | +| 还需要新的命令行参数（如 GDPO 的 `--gdpo-reward-keys`） | 4–6 个（上述 + `arguments.py` 的参数声明与校验 + 示例 + 文档） | + +注册表消除的是「同一个算法名散落在 6 处 if/elif」，不是「新增算法零成本」。GDPO 走的是最后一档。 + +`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`。适用于没有重要性比值修正的目标函数 | +| `supports_fully_async` | 设为 `False` 可拒绝 `--fully-async`（该模式下 advantage 由单副本服务按切片计算，无 DP 通信域） | +| `allows_reward_post_process_hooks` | 设为 `False` 可同时拦住 `--custom-reward-post-process-path` 与 `--agentic-custom-advantage-path`——这两个钩子都会在归一化器之前从 `post_process_rewards` 返回，静默跳过本算法的奖励阶段 | +| `uses_reward_components` | 算法消费多个具名奖励分量而非单个标量，驱动 `--gdpo-reward-keys` 校验 | + +表里除 `kl_level`、`needs_full_log_probs` 和 `advantage_normalization` 之外的字段，都由 `relax/utils/arguments.py` 的四个 `validate_*` 函数统一消费，**声明即生效**，不需要再去 `arguments.py` 加 `if`。（拆成四个是因为参数校验本身有推导顺序——例如 `--kl-coef` 必须在「检查 `--ref-load` 是否存在」之前判掉，one-update 等式必须在 `global_batch_size` 定稿之后判——与算法特殊性无关。）那三个字段是在 `relax/backends/megatron/loss.py` 里读的：新增一个前所未有的取值需要在那里加分支，复用已有取值则不用。 + +### 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 保持不变——多奖励算法（如 GDPO）也是在这一层把各分量收敛成一个标量的。 + +**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) +- GDPO 是最近一个走完整个流程的例子，可以对照 `relax/algorithms/` 与 `examples/gdpo/` 阅读。 diff --git a/docs/zh/guide/configuration.md b/docs/zh/guide/configuration.md index 9da9012ac..9059aaa4d 100644 --- a/docs/zh/guide/configuration.md +++ b/docs/zh/guide/configuration.md @@ -299,7 +299,7 @@ bash scripts/training/text/run-qwen3-4B-fp16-8xgpu.sh \ | 参数 | 类型 | 默认值 | 可选值 | 说明 | |------|------|--------|--------|------| -| `--advantage-estimator` | str | grpo | `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`、`gdpo`、`ppo`、`reinforce_plus_plus`、`reinforce_plus_plus_baseline` | 优势估计器。以 `--help` 为准：取值直接读注册表，新增算法无需改这张表即可出现。OPD 独立于该选项，使用 `--use-opd` 及对应 KL/loss 系数启用 | | `--normalize-advantages` | flag | False | - | 是否归一化优势 | | `--disable-grpo-std-normalization` | flag | - | - | 禁用 GRPO 标准差归一化（来自 [Dr.GRPO](https://arxiv.org/pdf/2503.20783)） | | `--disable-rewards-normalization` | flag | - | - | 禁用 reward 归一化 | diff --git a/examples/algorithms/README.md b/examples/algorithms/README.md index 3ee0f407b..45850353b 100644 --- a/examples/algorithms/README.md +++ b/examples/algorithms/README.md @@ -18,6 +18,7 @@ Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator | **CISPO** | `--advantage-estimator cispo` | 保留梯度方向、需要更高精度 | | **GSPO** | `--advantage-estimator gspo` | 序列级约束、稳定训练 | | **SAPO** | `--advantage-estimator sapo` | 平滑优化、soft 信任域 | +| **GDPO** | `--advantage-estimator gdpo` | 多奖励、分量独立归一化 | ## 选择建议 @@ -68,6 +69,13 @@ Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator - 梯度流更平滑，避免梯度突变 - 适合对稳定性要求高的场景 +### GDPO（多奖励） + +- 每个 reward 分量分别做组内标准化，再加权合并 +- 某个分量组内塌缩时，其它分量仍保留学习信号（GRPO 会丢掉整组） +- 适合 correctness + format、correctness + length 这类多目标任务 +- 详见 [examples/gdpo/](../gdpo/README.md) + ## 快速开始 ### 基础操作：修改算法参数 @@ -156,13 +164,13 @@ bash scripts/training/text/run-qwen3-4B-8xgpu.sh ### 通用参数 -| 参数 | 默认值 | 说明 | -| ----------------------- | -------------------- | --------------------------------------------------------------------------------------------------------------- | -| `--advantage-estimator` | `grpo` | 算法类型：`grpo`, `cispo`, `gspo`, `sapo`, `ppo`, `rloo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline` | -| `--eps-clip` | `0.2` | 下方裁剪边距（ratio 下界 = `1 - eps_clip`） | -| `--eps-clip-high` | 与 `--eps-clip` 相同 | 上方裁剪边距（ratio 上界 = `1 + eps_clip_high`） | -| `--clip-grad` | — | 梯度裁剪范数，CISPO 下推荐设为 `1.0` | -| `--kl-coef` | `0.0` | KL 惩罚系数；当前同步 PPO 会将非零值重置为 `0.0`，REINFORCE++ 等算法可使用 | +| 参数 | 默认值 | 说明 | +| ----------------------- | -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `--advantage-estimator` | `grpo` | 算法类型：`grpo`, `cispo`, `gspo`, `sapo`, `rloo`, `gdpo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`（取值由 `relax/algorithms/spec.py` 的注册表生成，以 `--help` 为准） | +| `--eps-clip` | `0.2` | 下方裁剪边距（ratio 下界 = `1 - eps_clip`） | +| `--eps-clip-high` | 与 `--eps-clip` 相同 | 上方裁剪边距（ratio 上界 = `1 + eps_clip_high`） | +| `--clip-grad` | — | 梯度裁剪范数，CISPO 下推荐设为 `1.0` | +| `--kl-coef` | `0.0` | KL 惩罚系数；当前同步 PPO 会将非零值重置为 `0.0`，REINFORCE++ 等算法可使用 | ### RLOO 专用约束与指标 @@ -194,6 +202,17 @@ bash scripts/training/text/run-qwen3-4B-8xgpu.sh | `--sapo-tau-pos` | `1.0` | Positive advantage 的温度参数 | | `--sapo-tau-neg` | `1.05` | Negative advantage 的温度参数（更高 = 更强抑制） | +### GDPO 专用参数 + +| 参数 | 默认值 | 说明 | +| ------------------------ | -------- | --------------------------------------------------------------- | +| `--gdpo-reward-keys` | — | **必填**，至少两个。奖励 dict 中要独立归一化的 key | +| `--gdpo-reward-weights` | 全 `1.0` | 各分量权重，长度须与 keys 一致；乘在**归一化后**的 advantage 上 | +| `--reward-key` | — | **必填**，选出 metrics 与 `raw_reward` 用的标量 | +| `--n-samples-per-prompt` | — | 必须 ≥ 2 | + +GDPO 与 `--normalize-advantages`、`--custom-reward-post-process-path`、`--agentic-custom-advantage-path` 和 `--fully-async` 互斥，参数校验阶段会报错。 + ### PPO 专用参数 | 参数 | 默认值 | 说明 | @@ -259,6 +278,7 @@ GSPO_ARGS=( examples/algorithms/ ├── README.md (本文件) ├── run-qwen35-9B-8xgpu-openr1mm-cispo-async.sh (CISPO 多模态示例) +├── ../gdpo/ (GDPO 双奖励示例) ├── ... (其他算法脚本) ``` @@ -272,6 +292,7 @@ examples/algorithms/ - **CISPO**：需要精细学习信号时更好，但需要 KL 约束 - **GSPO**：长序列任务，训练更稳定 - **PPO**：如果已有 Critic 资源，性能可能更好 +- **GDPO**：多个奖励分量各自需要归一化时用它 ### Q: CISPO 的梯度波动很大，正常吗？ @@ -293,3 +314,4 @@ examples/algorithms/ - [REINFORCE++ - Simple Efficient Alignment](https://arxiv.org/abs/2501.03262) - [RLOO - Back to Basics (Ahmadian et al. 2024)](https://arxiv.org/abs/2402.14740) - [RLOO - Buy 4 REINFORCE Samples, Get a Baseline for Free (Kool et al. 2019)](https://arxiv.org/abs/1905.12705) +- [GDPO - Group reward-Decoupled Normalization](https://arxiv.org/abs/2601.05242) diff --git a/examples/gdpo/README.md b/examples/gdpo/README.md new file mode 100644 index 000000000..283a8bba1 --- /dev/null +++ b/examples/gdpo/README.md @@ -0,0 +1,130 @@ +# GDPO 示例：correctness + format 双奖励 + +本目录是 GDPO（[arXiv 2601.05242](https://arxiv.org/abs/2601.05242)）的最小可运行示例，用 Qwen3-0.6B 单卡在 GSM8K 上训练。 + +## 为什么需要 GDPO + +奖励函数 `reward_gdpo.py` 返回两个分量： + +- `correctness` —— `` 标签里的答案是否正确（0 或 1） +- `format` —— 输出是否同时带 `` 与 ``（0、0.5 或 1） + +这两个分量会**不同步**：模型可能答对但没按格式输出，也可能格式完美但答错。 + +GRPO 把它们**加起来**再做一次组内归一化。这一步会丢掉两类信息。 + +**一、分量之间的相关结构。** GRPO 对和做组内标准化，每一组的输出恒为单位方差——无论这组的两个分量是彼此印证还是互相矛盾。GDPO 先让每个分量各自单位方差再加权求和，组合结果的方差是 + +``` +Var = Σ wᵢ² + 2 Σᵢ<ⱼ wᵢwⱼ ρᵢⱼ 等权重两分量时 = 2 + 2ρ +``` + +**由分量间的相关系数 ρ 决定**： + +| ρ | 含义 | 组合信号 | +| ---- | ---------------------- | ----------------- | +| → +1 | 两个分量指向同一批样本 | 增强（实测 2 倍） | +| 0 | 互不相关 | √2 倍 | +| → −1 | 两个分量互相矛盾 | 减弱，极限时归零 | + +第三步的 batch 白化是**跨组**的，所以这个组间强度差异会保留到最终 advantage。GRPO 看到的则是所有组都一样强。 + +**二、分量之间的尺度差异。** `correctness ∈ {0,1}` 与一个取值上百的 `format`（或论文实验里的响应长度）相加时，和的方差几乎全部来自大尺度那一维，GRPO 的方向就由它单独决定。GDPO 先让每个分量单位方差，权重才真正表达「相对重要性」而不是量纲。极端一点：`correctness=[1,1,0,0]`、`format=[0,100,200,300]`（两者排序相反）时，GRPO 给答错的长响应最高 advantage，GDPO 不会。 + +## GDPO 不能做什么 + +**各分量在组内恰好加和为常数时，GDPO 也救不回来。** 这正是上表 `ρ = −1` 那一行：若 `correctness + format ≡ C`，则 `format = C − correctness`，两者标准化后恒有 `z_format = −z_correctness`，**等权重下完全抵消为零**——与 GRPO 得到同样的结果。 + +这一点本文档此前写反了，说这正是 GDPO 的优势场景。它不是：那是 `ρ = −1`，是组合方差 `2 + 2ρ` 恰好取零的极端。只有不等权重（如 `--gdpo-reward-weights 2.0 1.0`）能打破这个平衡。 + +（另外，若两个分量在组内**都**恒定，GDPO 与 GRPO 一样返回零，不会无中生有。） + +## 运行 + +```bash +export MODEL_DIR=/path/to/models # 需含 Qwen3-0.6B +export DATA_DIR=/path/to/data # 需含 gsm8k/train.jsonl +export EXP_DIR=/path/to/experiments + +bash examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh +``` + +### 数据要求 + +每条 prompt 需要**自带**格式要求，否则基座模型不会产出 ``/`` 标签，`format` 分量会恒为 0（组内塌缩），GDPO 就退化成只看 `correctness`。准备数据时给 question 追加一句即可： + +```python +instruction = ( + "\n\nThink step by step inside tags, then give only the " + "final number inside tags." +) +df["question"] = df["question"] + instruction +``` + +**不要用 `--system-prompt` 代替**：`relax/utils/data/data_utils.py:181` 把 system message 的 content 构造成多模态 list（`content: [{"type": "text", ...}]`），Qwen3-0.6B 这类纯文本 chat template 渲染时会报 +`TypeError: can only concatenate str (not "list") to str`。这是既有的框架限制，与 GDPO 无关。 + +## 参数说明 + +```bash +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format # 独立归一化的分量，至少两个 + --gdpo-reward-weights 1.0 1.0 # 可省略，默认全 1 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score # 必填：metrics 与 raw_reward 用的标量 + --n-samples-per-prompt 8 # 必须 ≥ 2 +) +``` + +`--gdpo-reward-weights` 乘的是**归一化之后**的 advantage，不是原始 reward。经过第一步各分量已经是单位方差，所以权重表达的是相对重要性，而不是分量的量纲——把 `format` 的取值范围从 `[0,1]` 改成 `[0,100]` 基本不改变训练结果。 + +说「基本」而不是「完全」，是因为第一步除的是 `std + 1e-4` 而不是 `std`。当某分量的组内标准差本身就落到 `1e-4` 量级时，这个加性 epsilon 会随量纲变化而改变阻尼比例，缩放就不再是严格等价的。二值奖励（std≈0.4）离这个区间很远，连续奖励（如响应长度）则可能撞上。 + +## 换成自己的奖励 + +改 `reward_gdpo.py` 的 `compute_gdpo_reward`，返回的 dict 需要包含 `--gdpo-reward-keys` 列出的全部 key，外加 `--reward-key` 指定的那个标量。 + +分量缺失、非数值、bool、NaN/Inf 都会直接报错。这是有意的：静默填 0 会让一个坏掉的奖励函数看起来像是「这一维恰好塌缩了」，训练照跑，问题要很久以后才暴露。 + +## 已知偏差 + +1. **第三步的 batch 边界（已正确处理）**。调用方为效率会先合并多个训练批再调用 advantage，但 `_whiten_by_segment` 用 `mini_batch_sizes` 把它们切回**每个 optimizer 训练批各自白化**，因此 `num_rollout_minis > 1` 时仍对齐论文 Eq. 6，**不要求** `rollout_batch_size × n_samples_per_prompt == global_batch_size`。本脚本把 `4 × 8` 与 `--global-batch-size 32` 设成相等只是让例子最简单，并非必需——把 `--global-batch-size` 改成 16 就会跑出两段。跨 DP 的 all-reduce 保证统计量覆盖全部 rank。**`--fully-async` 会在参数校验阶段被拒绝**——那条路径的切片可能小到只有一个样本，白化输出恒为 0。 + + 调用方**必须**提供这份切分（`rollout_mini_local_sample_counts`）。缺失时 GDPO 直接报错，不回退到「整个 rollout 白化一次」：那不是同一个目标的粗糙版本，而是另一个目标——`test_merging_the_batches_would_flip_signs_not_just_rescale` 里 8 个样本有 4 个符号翻转，而 loss、grad_norm、advantage 均值全都正常。 + +2. **单个奖励时 GDPO 不等于 GRPO**。step1 除以 `std_g + 1e-4`、GRPO 除以 `std_g + 1e-6`，各组 `std_g` 不同 → 尺度因子逐组不同，step3 还会再做一次 batch 白化，所以不是「差一个正标量」那么简单。要 GRPO 语义就用 `--advantage-estimator grpo`。 + +3. **`--n-samples-per-prompt 2` 时幅度信息丢失**：任意两个不同值标准化后恒为 ±0.7071。示例用 8 就是为了避开这一点。 + +4. **恒和分量的数值假信号（未解决，如实记录）**。若两个分量恰好满足 `r₂ = C − r₁`，数学上应完全抵消为零（见上文「GDPO 不能做什么」）。这里的标准化在这种输入下是**病态**的：它要除以一个接近零的 std，相对误差被放大 `max|x| / std` 倍。 + + 奖励曾经在进入归一化前被 cast 成 float32，而 `C − r` 这种值不一定能被 float32 精确表示——两者相加仍舍回 `C`（看起来恒和），各自却已偏离。残差量级约 `1e-1`，**比信号本身还大**，第三步再除以同量级的 batch 标准差，输出就是 O(1) 的 advantage：一个本无信号的组拿到了方向由舍入决定的梯度。实测一组和为 `308.95172119140625` 的奖励，最终 advantage 为 `[-0.5770, 1.1539, -0.5770]`。 + + **分量全链路改成 float64**（`extract_reward_components`）之后，这个特定输入的残差降到 `1e-10` 量级，第三步分母里的 `GDPO_EPS = 1e-4` 钳住放大倍数，输出降到 `1e-6`。但**问题没有被解决，只是被推远了**：残差大致是 `ulp(C) / 组内展布`，随基数增长。同样的构造在 `C = 1e13`、或两列都很大且跨 binade 时，残差的 std 超过 `GDPO_EPS`，钳制失效，输出回到 **O(1)**。 + + 曾经有两个机制试图挡住它，都已移除： + + - **`combine_group` 里的噪声地板**（按 `8 · Σ|wₖ|·noiseₖ` 整组置零）。它摧毁真实信号：两分量 `base = 4.05e13` 时地板 0.148 而真实信号 0.033；地板对分量求和而任何逐列噪声度量取 max，16 个各自干净的分量能把它推到 0.82。而且它修改的是训练值本身。 + - **一个相对幅度判据**（`|Σwz|` 相对 `Σ|wₖ|·max|zₖ|` 低六个数量级即判为残差）。它错得更根本：分子正比于权重之**差**、分母正比于权重的**大小**，所以它测的是权重配置而非数据；`G = 2` 时精确退化为 `|w₁−w₂|/(|w₁|+|w₂|)`，与任何奖励值无关。实测它在两个方向同时判反——扔掉一个最终 advantage 0.43 的真信号，放行一个 1.08 的纯舍入结果。这不是阈值问题：`G ≥ 3` 时中心化子空间至少二维，可构造 `z₂ = -z₁ + δu`（`u ⊥ z₁`，δ 任意小），所以真实信号的比值没有正下界。 + + 所以现在的状态是：**`combine_group` 严格返回 Eq. 7，恒和组带着它的舍入残差进训练，没有任何机制识别它**。实测到达 optimizer 的量级（`[C−δ, δ]`，δ∈{0.1,0.2,0.3,0.7}，等权）： + + | C | 自己独占一个白化单元 | 与 8 个健康组共享一个 | + | ---- | -------------------- | --------------------- | + | 1e8 | 2.3e-4 | 2.7e-8 | + | 1e9 | 2.6e-3 | 3.1e-7 | + | 1e11 | 2.0e-1 | 3.3e-5 | + | 1e13 | 1.2 | 5.2e-3 | + + **第二列要连着 C 一起读。** 本文档早先的版本只引了它 `C=1e9` 那一格、写成「混批约 1e-7」，读起来像个普遍上界——不是。与健康组共享白化单元只是把残差除以**它们的**标准差（这里是几百倍的常数因子），不改变它随 C 的增长；C=1e13 时共享单元里仍有 5e-3。 + + 而且「白化单元」不等于「整个 rollout」：`_whiten_by_segment` 按**训练批**分别白化，所以只有当健康组落进同一个训练批时才有这个除法。退化组独占一段时，直接回到第一列。没有任何机制保证这种混合，它取决于 rollout 怎么被切分。 + + GRPO 没有这个问题：它归一化的是**和**，而和在 float32 下确实恒定，塌缩检测会直接置零。 + +## 冲突项 + +`--normalize-advantages`、`--custom-reward-post-process-path`、`--agentic-custom-advantage-path` 和 `--fully-async` 都不能与 GDPO 同用，参数校验阶段会直接报错。第一个会造成双重白化；第二、三个都会赶在归一化器之前从 `post_process_rewards` 返回，导致 GDPO 的前两步被静默跳过（这两个由 `AlgorithmSpec.allows_reward_post_process_hooks` 一起把守）；第四个的统计窗口可能小到只剩一个样本。 + +配 `--dynamic-sampling-filter-path` 时**不冲突**：内置的 `check_reward_nonzero_std` 会算出 GDPO 前两步的组合结果，再在**训练实际使用的 float32** 上判 `min != max`——与单奖励算法那条分支问的是同一个问题，只是维度从 1 变成 K。判据不含任何容差：权重把某分量静音、或两个分量精确抵消，组合结果就是精确的零，能判出来；而近似抵消判不出来，那种组会带着上面说的舍入残差进训练。只有指向**自定义** filter 时才会警告——那种 filter 若只看 `--reward-key` 标量，就会丢掉只在其它分量里有信号的组。 diff --git a/examples/gdpo/__init__.py b/examples/gdpo/__init__.py new file mode 100644 index 000000000..9f3863608 --- /dev/null +++ b/examples/gdpo/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. diff --git a/examples/gdpo/reward_gdpo.py b/examples/gdpo/reward_gdpo.py new file mode 100644 index 000000000..f0e694792 --- /dev/null +++ b/examples/gdpo/reward_gdpo.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Two-component reward for the GDPO example: correctness and format. + +GDPO standardizes each component within its prompt group before combining them. +Two things survive that GRPO's single standardization of the summed reward +destroys: + +* **Relative strength between groups.** Standardizing per group forces every + group to unit variance, so a group where only `correctness` varies and one + where both components vary come out identical. Standardizing each component + first makes the latter twice the amplitude, and step 3 whitens across the + *batch*, so the difference reaches the final advantage. +* **Scale disparity between components.** A `correctness` in {0, 1} added to a + reward in the hundreds gives a sum whose variance is essentially the large + component's, so GRPO's direction is decided by it alone. + +What GDPO does *not* do: rescue a group whose components sum to a constant. +There `format = C - correctness` forces the standardized values to be exact +opposites, and equal weights cancel them to zero -- the same answer GRPO +gives. Only unequal weights break that tie. (If every component is constant +within the group, GDPO returns zero too.) + +Wire it up with:: + + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + +``score`` is what ``--reward-key`` selects for metrics and for the ``raw_reward`` +column; the training signal comes from the two components. +""" + +import re +from typing import Any + + +_ANSWER_RE = re.compile(r"(.*?)", re.DOTALL) +_THINK_RE = re.compile(r".*?", re.DOTALL) + + +def _extract_answer(response: str) -> str | None: + match = _ANSWER_RE.search(response) + return match.group(1).strip() if match else None + + +def _final_answer(label: Any) -> str: + """Normalise a GSM8K label to just the final answer. + + GSM8K ships ``answer`` as the full worked solution ending in ``#### 36``, so + comparing a model's ``36`` against the whole string makes + ``correctness`` zero for every rollout. That collapses the component in every + group, and GDPO silently degrades to the single ``format`` reward -- which is + the one thing this example exists to demonstrate it does not do. + + Handling it here rather than only in a data-prep step keeps the example + working against the dataset it names. Labels without the marker pass through + unchanged, so a pre-cleaned dataset behaves identically. + """ + text = str(label).strip() + return text.rsplit("####", 1)[-1].strip() if "####" in text else text + + +def compute_gdpo_reward(response: str, label: Any) -> dict[str, float]: + """Score one response on answer correctness and on output format. + + The two components are deliberately decorrelated: a response can be correct + without the expected tags, and well-formatted while wrong. + + ``score`` is ``correctness`` rather than the sum, on purpose. It feeds + ``--reward-key``, which selects the scalar for metrics and the + ``raw_reward`` column only -- it does not participate in the GDPO + computation, which reads the two components directly. Reporting accuracy + there is more legible on a dashboard than a blended number. Note this means + ``rollout/raw_reward`` tracks correctness alone, not overall reward. + """ + answer = _extract_answer(response) + correctness = 1.0 if answer is not None and answer == _final_answer(label) else 0.0 + + has_think = _THINK_RE.search(response) is not None + format_score = 0.5 * float(has_think) + 0.5 * float(answer is not None) + + return { + "score": correctness, + "correctness": correctness, + "format": format_score, + } + + +async def reward_func(args: Any, sample: Any, **kwargs: Any) -> dict[str, float]: + """Entry point for ``--custom-rm-path``.""" + return compute_gdpo_reward(sample.response, sample.label) diff --git a/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh new file mode 100644 index 000000000..66856dc2d --- /dev/null +++ b/examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh @@ -0,0 +1,143 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3-0.6B single-GPU GDPO training on GSM8K. +# +# GDPO (arXiv 2601.05242) standardizes each reward component within its prompt +# group before combining them. The combined advantage is sum_k w_k * z_k with +# each z_k at unit variance, so its variance is 2 + 2*rho for two equal +# weights: what GDPO preserves is the CORRELATION between components, which +# GRPO destroys by standardizing the summed reward to unit variance for every +# group regardless. Corroborating components amplify, contradicting ones +# attenuate, and at rho = -1 they cancel -- so it does NOT rescue a group whose +# components sum to a constant. Separately and unconditionally, GRPO's +# direction is dominated by whichever component has the largest raw spread; +# here each contributes according to its weight. See examples/gdpo/README.md. +# The reward function in reward_gdpo.py returns both components; +# --gdpo-reward-keys names them. +# +# Usage: +# bash examples/gdpo/run-qwen3-0.6B-1xgpu-gdpo.sh + +set -ex +set -o pipefail + +now=$(date "+%Y-%m-%d-%H:%M:%S") + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# Auto-source local environment when not launched via an external entrypoint +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../scripts/entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" + +PROJECT_NAME="${PROJECT_NAME:=relax-gdpo}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" +NUM_ROLLOUT="${NUM_ROLLOUT:=20}" + +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3-0.6B + --ref-load ${MODEL_DIR}/Qwen3-0.6B + --load ${EXP_DIR}/Qwen3-0.6B_mcore_gdpo/ + --save ${EXP_DIR}/Qwen3-0.6B_mcore_gdpo/ + --save-interval 100 + --max-actor-ckpt-to-keep 1 + --megatron-to-hf-mode bridge +) + +# NOTE: the format instruction belongs in the prompt text, not in +# --system-prompt. relax/utils/data/data_utils.py:181 builds the system message +# with multimodal list content (`content: [{"type": "text", ...}]`), which a +# text-only chat template such as Qwen3-0.6B's cannot render: +# TypeError: can only concatenate str (not "list") to str +# Prepare the dataset so each question already asks for the / +# tags; see examples/gdpo/README.md. +ROLLOUT_ARGS=( + --prompt-data ${DATA_DIR}/gsm8k/train.jsonl + --input-key question + --label-key answer + --apply-chat-template + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 4 + --n-samples-per-prompt 8 + --rollout-max-response-len 1024 + --rollout-temperature 1.0 + # 4 * 8 == 32 just keeps this example minimal. GDPO's step 3 stays aligned + # with Eq. 6 regardless: _whiten_by_segment splits any merged rollout back + # into per-optimizer-batch segments and whitens each on its own, so + # num_rollout_minis > 1 does NOT require rollout_batch_size * n == gbs. + --global-batch-size 32 +) + +# The two reward components come from examples/gdpo/reward_gdpo.py. +# --reward-key selects the scalar used for metrics and the raw_reward column; +# --gdpo-reward-keys names the components GDPO standardizes independently. +GDPO_ARGS=( + --advantage-estimator gdpo + --gdpo-reward-keys correctness format + --gdpo-reward-weights 1.0 1.0 + --custom-rm-path examples.gdpo.reward_gdpo.reward_func + --reward-key score + --eps-clip 0.2 + --kl-coef 0.00 + --entropy-coef 0.00 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --clip-grad 1.0 +) + +WANDB_ARGS=( + --use-tensorboard + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name qwen3-0.6b-gdpo-gpu1-${now} +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.5 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +mkdir -p log +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 1], "rollout": [1, 1]}' \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --use-health-check \ + --balance-data \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GDPO_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen3-0.6b-gdpo-gpu1-${now}.log diff --git a/examples/generate_reward_model/post_process_genrm_swap.py b/examples/generate_reward_model/post_process_genrm_swap.py index b3dcf4697..8a61290f1 100644 --- a/examples/generate_reward_model/post_process_genrm_swap.py +++ b/examples/generate_reward_model/post_process_genrm_swap.py @@ -100,19 +100,29 @@ 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 is also true + of ``gdpo_decoupled``, which standardises several reward components + independently and is not what the code below computes. GDPO cannot reach + here today (``allows_reward_post_process_hooks=False`` rejects this hook in + ``relax/utils/arguments.py``), so this is a guard against a future + multi-reward algorithm that allows the hook, not a live bug. Naming + normalizers keeps it registry-driven: a new algorithm reusing either one is + covered for free, and a genuinely new normalizer is exactly the case where + this reimplementation needs a human to look at it. + """ + from relax.algorithms import get_algorithm + + 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 +131,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/agentic/rollout.py b/relax/agentic/rollout.py index 622fa5eb7..e0cdad0c8 100644 --- a/relax/agentic/rollout.py +++ b/relax/agentic/rollout.py @@ -19,6 +19,8 @@ get_agentic_runtime_resources, ) from relax.agentic.profile import TRACE_KEY +from relax.algorithms import get_algorithm +from relax.algorithms.rewards import metrics_group_verdict, zero_std_group_label from relax.engine.filters.base_types import MetricGatherer, call_dynamic_filter from relax.engine.rollout import on_policy_distillation as opd from relax.engine.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput @@ -1293,19 +1295,30 @@ def _dict_add_prefix(d, prefix): def _compute_zero_std_metrics(args, all_samples: list[Sample]) -> dict[str, float]: - if args.advantage_estimator == "ppo": + if get_algorithm(args.advantage_estimator).needs_critic: return {} all_sample_groups = group_by(all_samples, lambda sample: sample.group_index) - reward_groups = [ - [sample.get_reward_value(args) for sample in group if sample.reward is not None] - for group in all_sample_groups.values() - ] - interesting_rewards = [ - str(round(rewards[0], 1)) - for rewards in reward_groups - if rewards and all(rewards[0] == reward for reward in rewards) - ] + interesting_rewards = [] + for group in all_sample_groups.values(): + # Counted as flat only when it carries no signal *for this algorithm*: + # for a multi-reward one that means the weighted combination cancels, + # not that the --reward-key scalar happens to be flat. + # + # `is not True`, so a group that is unscored or unreadable is skipped + # rather than counted -- which is what this metric did before, since it + # dropped `reward is None` samples and then required a non-empty list. + # The distributed copy counts them instead; see `metrics_group_verdict`. + if metrics_group_verdict(args, group) is not True: + continue + # A `True` verdict already implies a scored sample, so the label is + # never None here. It goes through the shared helper anyway: this line + # and its opposite number in the distributed copy are the half of this + # metric `metrics_group_verdict` did not cover, and they had already + # drifted -- that copy read the label off `group[0]`, scored or not. + label = zero_std_group_label(args, group) + if label is not None: + interesting_rewards.append(label) return {f"zero_std/count_{reward}": len(items) for reward, items in group_by(interesting_rewards).items()} diff --git a/relax/algorithms/__init__.py b/relax/algorithms/__init__.py new file mode 100644 index 000000000..2d8c798bf --- /dev/null +++ b/relax/algorithms/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Algorithm registry: names, capabilities and implementations in one place. + +Nothing in this package may import ``megatron``, ``ray``, ``transfer_queue``, +``tensordict``, ``relax.components`` or ``relax.backends`` at module level. The +registry is imported by argument parsing and by both worker processes, so a +heavy import here would pull the whole training stack into ``--help`` and into +the CPU-only test runner. +""" + +from relax.algorithms.spec import ALGORITHM_SPECS, AlgorithmSpec, get_algorithm, list_algorithm_names + + +__all__ = ["ALGORITHM_SPECS", "AlgorithmSpec", "get_algorithm", "list_algorithm_names"] diff --git a/relax/algorithms/advantages.py b/relax/algorithms/advantages.py new file mode 100644 index 000000000..fe8e128fd --- /dev/null +++ b/relax/algorithms/advantages.py @@ -0,0 +1,437 @@ +# 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 +import torch.distributed as dist + +from relax.algorithms.numerics import GDPO_EPS, any_rank_has_non_finite, distributed_mean_std, is_collapsed +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 whiten_scalar(values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None) -> torch.Tensor: + """Sequence-level whitening of one scalar per sample. + + This is GDPO's batch-wise normalisation (arXiv 2601.05242, Eq. 6). It is + deliberately *not* the token-level ``distributed_masked_whiten`` used by + ``--normalize-advantages``: weighting by token count would let long + responses dominate the statistics, which Eq. 6 does not do. + + ``process_group`` must be supplied wherever the caller holds only a shard of + the values. In the Megatron path each data-parallel rank owns + ``num_rollout_minis * global_batch_size / dp_size`` samples — its shard of + the whole rollout, merged before ``compute_advantages_and_returns`` runs — so + whitening locally would give every rank its own mean and scale, not the "one + global scale factor" the maths assumes. ``None`` means "I own every value you + need"; note the single-replica ``Advantages`` deployment cannot be that caller + for GDPO, since ``supports_fully_async=False`` rejects the configuration that + would route here — so in practice ``None`` is the CPU-side and test callers. + + Scope of one call: this whitens exactly the tensor it is handed. When the + caller has merged several training batches, :func:`_whiten_by_segment` splits + them back and calls this once per optimizer batch, so each call stays aligned + with Eq. 6's per-batch statistic regardless of ``num_rollout_minis``. Only + the un-segmented path (``mini_batch_sizes=None``) whitens a whole merged + tensor at once. + + A batch where every value is identical returns exact zeros; see + :func:`relax.algorithms.numerics.is_collapsed`. + + A batch containing a non-finite value raises. It used to return zeros, on + the reasoning that a non-finite standard deviation is "no spread" -- but + "every sample scored the same" and "the arithmetic overflowed" are opposite + diagnoses, and the second one is reachable: the reward stage checks + finiteness in float64 while the values arrive here in float32 + (``dict_to_tensordict`` casts the rewards column, and ``_as_reward_tensor`` + casts again), so a combined reward of 5.999e38 is finite where it was + checked and ``inf`` where it is used. + + **The whitening itself runs in float64.** That is not a precision nicety; + it closes two holes a finiteness check on the input cannot: + + * ``distributed_mean_std`` accumulates in float64 but returns + ``std.to(values.dtype)``. Two finite float32 values, ``[-FMAX, FMAX]``, + give an unbiased std of ``sqrt(2) * FMAX`` -- finite in float64, ``inf`` + once cast back to float32. Dividing by that returned an all-zero batch, + silently. Handing this function's float64 copy to ``distributed_mean_std`` + means the narrowing cast never happens. + * ``values - mean`` was itself a float32 subtraction. ``[FMAX] * 10 + + [-FMAX]`` has a finite mean *and* a finite std, and the subtraction still + overflowed to ``-inf`` for the last element. Restoring a guard on ``std`` + alone would not have caught that one. + + An earlier version of this comment argued the ``std`` guard could not fire + because float64 accumulation needs ~1e153 samples to overflow. Both halves + were wrong: the arithmetic is ~1.5e231, and the overflow does not come from + accumulation at all -- it comes from the two narrowing casts above, and two + samples are enough. + + The result is normalised, so it is bounded by the batch's own spread and + casts back to the caller's dtype without overflowing. + """ + # Collective, not a local `isfinite`. The `is_collapsed` reduction is two + # lines below, so a rank that raised here on its own would hang every other + # rank in the group -- the deadlock `_agree_on_segmentation` was written to + # prevent, reintroduced by a check that looks purely local. + if any_rank_has_non_finite(values, process_group=process_group): + raise ValueError( + "GDPO batch whitening received a non-finite advantage. The reward stage verified these " + "in float64, so this is the float32 hand-off overflowing: the combined reward fits in " + "float64 but not in the dtype it is carried and trained in. Rescale the reward " + "components or their --gdpo-reward-weights." + ) + # Collapse is judged on the values as the trainer will see them, not on the + # float64 copy: a batch differing only below float32's resolution trains on + # nothing, so it counts as collapsed here even though float64 can still tell + # the values apart. + if not values.is_floating_point(): + # `.to(values.dtype)` at the end would truncate every z-score to the + # integer 0 and hand back a silently all-zero batch. The previous + # version returned floats here by accident of the division; neither is + # a contract worth having, so say so instead. + raise ValueError(f"GDPO batch whitening needs a floating-point tensor, got {values.dtype}.") + if is_collapsed(values, process_group=process_group): + return torch.zeros_like(values) + work = values.double() + mean, std = distributed_mean_std(work, process_group=process_group) + # Stacked rather than `isfinite(mean) and isfinite(std)`: Python's `and` + # calls `__bool__` on the first tensor and, when it is true, on the second, + # so the readable spelling costs one or two device-to-host syncs per + # training batch. One stack, one reduction, one read. + if not torch.isfinite(torch.stack((mean, std))).all(): + # Unreachable from float32 -- the largest float32 squared is 1.2e77 + # against float64's 1.8e308 -- but reachable from a float64 caller near + # its own maximum, where the squares overflow inside the reduction. + # Left unchecked that is the all-zero batch again, one dtype up. + raise ValueError( + f"GDPO batch whitening overflowed computing {values.dtype} statistics in float64. " + "Rescale the reward components or their --gdpo-reward-weights." + ) + return ((work - mean) / (std + GDPO_EPS)).to(values.dtype) + + +def _as_reward_tensor(rewards: Any, kl: list[torch.Tensor]) -> torch.Tensor: + """Rewards as a detached float32 tensor on the KL tensors' device. + + ``detach()`` is defensive, and worth being precise about because the + comment here used to call it load-bearing, which overstates it. + + Both production call sites pass ``list[float]`` -- ``loss.py`` reads + ``rollout_data["rewards"]`` and the Advantages deployment reads the same + column off the TransferQueue -- so the tensor branch below is not currently + reached by anything, and removing the detach would change no observable + behaviour today. + + It stays because the refactor did quietly widen what this accepts. The + pre-registry code built the tensor with ``torch.tensor(rewards, ...)``, + which copies and drops autograd history even when handed a tensor; + ``.to()`` returns the *same* object when dtype and device already match. So + a future caller passing a reward tensor with ``requires_grad=True`` would, + without the detach, get advantages carrying grad history into the policy + loss. Rewards are data, not something to backprop into. + """ + if isinstance(rewards, torch.Tensor): + return rewards.detach().to(dtype=torch.float32, device=kl[0].device) + 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 _agree_on_segmentation(n_segments: int, local_error: str, values, process_group) -> None: + """Reach one verdict on the segmentation across the whole group. + + Each segment costs one collective, so this has to agree *before* any of + them run. Two ways it can disagree, both of which deadlock rather than + fail: + + * The segment counts differ. A rank expecting two segments and a rank + expecting three both block on the third collective -- no traceback, no + exit code, just a job that never finishes. ``num_rollout_minis`` comes + from the minibatch plan and is expected to be uniform, but "expected" is + not "checked". + * One rank's ``mini_batch_sizes`` is malformed. Raising locally is worse + than not checking at all: that rank leaves the collective sequence while + every other rank is still waiting inside it. + + So the local verdict travels *with* the count, in one MAX reduction (the + low bound negated, the way :func:`~relax.algorithms.numerics.is_collapsed` + does it). Every rank reads the same three numbers and therefore raises + together, at the same call, with a message naming which failure it was. + """ + if process_group is None: + if local_error: + raise ValueError(local_error) + return + + # A rank that already failed contributes 0 segments: a neutral value that + # cannot be mistaken for a real count, and one that trips the mismatch + # branch too if the `any_bad` branch were ever removed. + flags = torch.tensor( + [-n_segments if not local_error else 0, n_segments if not local_error else 0, 1 if local_error else 0], + dtype=torch.int64, + device=values.device, + ) + dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=process_group) + low, high, any_bad = -int(flags[0]), int(flags[1]), int(flags[2]) + + if any_bad: + raise ValueError( + local_error + or "another rank reported malformed mini_batch_sizes; every rank fails here so the " + "per-segment collectives cannot deadlock." + ) + if low != high: + raise ValueError( + f"mini_batch_sizes describes {n_segments} segment(s) on this rank, but ranks in the group " + f"report between {low} and {high}. Every rank must whiten the same number of segments, " + "or the per-segment collectives deadlock." + ) + + +def _whiten_by_segment(values, mini_batch_sizes, process_group): + """Whiten each training batch separately, in the order they were merged. + + ``num_rollout_minis`` comes from the minibatch plan rather than from the + data, so every rank is expected to run the same number of segments; + :func:`_agree_on_segmentation` is what turns that expectation into a + checked precondition rather than a deadlock. + + What it checks is the segment *count*, not segment *identity*. Two further + properties are relied on and not verified here: + + * Segment ``k`` holds training batch ``k`` on every rank. ``actor.py`` + fetches batches in ``batch_index`` order and appends the counts in that + same order, and the fetch is *addressed* rather than popped: + ``_get_data_from_transfer_queue`` passes ``batch_index`` to the + TransferQueue sampler, which keys its replay cache on + ``(partition_id, task_name, dp_rank, batch_index)``. Two ranks asking for + the same ``batch_index`` therefore receive their own shards of the *same + logical mini-batch* -- that addressing is the guarantee, not the fact + that both loops happen to count upwards. Nothing in this function would + notice if it stopped holding, and the failure would be silent -- + statistics mixed across two training batches, no error. + * A segment's sample count may differ between ranks (a data-parallel split + balances tokens, not samples) and that is fine, because the statistic is + reduced across the group. What is not fine is two ranks disagreeing about + *which* batch a segment belongs to. + + Both are guaranteed upstream rather than here because checking identity + would need a batch id in the metadata, which the plan does not currently + carry. If that metadata ever appears, fold it into the same MAX reduction. + """ + # Validate whatever was passed, including a single segment: treating an empty + # or malformed list as "fall back to one window" would silently restore the + # merged behaviour this function exists to replace. The verdict is not acted + # on until the whole group has shared it. + local_error = "" + if mini_batch_sizes is None: + # Not a fallback. Whitening the merged rollout as one window is a + # different optimisation target, not a coarser version of the same one: + # `test_merging_the_batches_would_flip_signs_not_just_rescale` shows + # half the advantages changing sign. Every Megatron producer writes + # these counts today, so a missing one means a new or changed caller, + # and the only thing worse than that caller failing is that caller + # training on the wrong objective while every metric stays finite. + n_segments = 0 + local_error = ( + "mini_batch_sizes is None: the caller did not supply " + "`rollout_mini_local_sample_counts`. GDPO whitens per training batch, so there is no " + "safe default -- whitening the merged rollout instead optimises a different objective." + ) + elif not mini_batch_sizes or any(not isinstance(n, int) or n <= 0 for n in mini_batch_sizes): + n_segments = 0 + local_error = f"mini_batch_sizes must be a non-empty list of positive ints, got {mini_batch_sizes}." + elif sum(mini_batch_sizes) != values.numel(): + n_segments = 0 + local_error = ( + f"mini_batch_sizes {mini_batch_sizes} sum to {sum(mini_batch_sizes)}, " + f"but this rank holds {values.numel()} samples." + ) + else: + n_segments = len(mini_batch_sizes) + + _agree_on_segmentation(n_segments, local_error, values, process_group) + + if n_segments == 1: + return whiten_scalar(values, process_group=process_group) + out, start = [], 0 + for size in mini_batch_sizes: + out.append(whiten_scalar(values[start : start + size], process_group=process_group)) + start += size + return torch.cat(out) + + +def advantage_gdpo(args: Any, *, rewards, kl, process_group=None, mini_batch_sizes=None, **_unused): + """GDPO step 3: whiten the combined per-sample advantage, then broadcast. + + Steps 1 and 2 (per-reward group standardisation and the weighted sum) ran + on the rollout side, so ``rewards`` already holds one ``A_sum`` per sample. + Doing step 3 here rather than there keeps it out of reach of ``--custom- + reward-post-process-path``, which short-circuits reward post-processing + entirely, and off the streaming transfer-batch boundary with its undersized + tail flushes. + """ + reward_tensor = _as_reward_tensor(rewards, kl) + whitened = _whiten_by_segment(reward_tensor, mini_batch_sizes, process_group) + returns = get_grpo_returns(whitened, kl) + advantages = list(returns) + return advantages, returns + + +def advantage_reinforce_plus_plus(args: Any, *, rewards, kl, loss_masks, response_lengths, total_lengths, **_unused): + """Discounted returns for REINFORCE++ + (https://arxiv.org/pdf/2501.03262).""" + 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, + "gdpo": advantage_gdpo, + "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, + process_group: dist.ProcessGroup | None = None, + mini_batch_sizes: list[int] | None = None, +) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Dispatch to the estimator registered for ``args.advantage_estimator``. + + ``mini_batch_sizes`` is this rank's per-training-batch sample counts, in the + order the caller merged them. Only estimators whose statistics are defined + per batch read it; the rest absorb it in ``**_unused`` and are unaffected. + + ``process_group`` is the group across which the batch is sharded, or + ``None`` when the caller holds every sample. Estimators that compute batch- + level statistics need it to see the whole batch. + + ``padded_total_lengths`` is likewise call-site specific: only the Megatron + path can compute it, and only GAE consumes it. Every parameter here is + keyword-only and every estimator absorbs the rest, so a call site that + 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, + process_group=process_group, + mini_batch_sizes=mini_batch_sizes, + ) diff --git a/relax/algorithms/numerics.py b/relax/algorithms/numerics.py new file mode 100644 index 000000000..e43d68f0d --- /dev/null +++ b/relax/algorithms/numerics.py @@ -0,0 +1,221 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Shared numerical constants and guards for the algorithm implementations. + +Both the reward stage and the advantage stage standardise values by dividing by +a standard deviation, so they need the same epsilon and the same notion of +"this group carries no signal". Keeping those here prevents the two stages +from drifting apart. +""" + +import torch +import torch.distributed as dist + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +_LOGGED_GROUP_SIZE = False + + +def _log_group_once(process_group: dist.ProcessGroup | None) -> None: + """Report the reduction group once per process. + + Whether a batch statistic is global or per-shard is invisible in the loss + curve: a misconfigured run where tensor parallelism ate the extra GPUs + leaves the data-parallel group at size 1, the all-reduce becomes an + identity, and everything still trains. This line is what makes that + distinguishable in a log. + """ + global _LOGGED_GROUP_SIZE + if _LOGGED_GROUP_SIZE: + return + _LOGGED_GROUP_SIZE = True + if process_group is None: + logger.info("Batch statistics are local (no process group); the caller owns the whole batch.") + else: + logger.info( + "Batch statistics reduce over dp_world=%d (this rank is dp_rank=%d).", + dist.get_world_size(process_group), + dist.get_rank(process_group), + ) + + +STD_EPS = 1e-6 +"""Epsilon added to a standard deviation before dividing by it. + +Matches the value the pre-registry GRPO path used, so GRPO, GSPO, SAPO and +CISPO keep producing exactly the numbers they produced before the registry +existed. That parity is the whole point of the equivalence tests, so this +constant is not free to move. +""" + +GDPO_EPS = 1e-4 +"""Epsilon for GDPO's two standardisation steps. + +Deliberately not :data:`STD_EPS`. The reference implementation +(``trl/trainer/grpo_trainer.py``, the ``scale_rewards`` GDPO path) divides by +``std + 1e-4`` at both the per-reward group step and the batch step, and GDPO +is new here, so there is no prior Relax behaviour that matching it would break. + +The choice only bites near-degenerate groups: with binary rewards and eight +samples the group std is around 0.4 and the two constants differ by 0.02%, but +a continuous reward (the paper's maths setup uses response length) can leave a +group with std ~1e-3, where 1e-6 and 1e-4 disagree by roughly 10% on the scale +factor. Exactly-collapsed groups never reach either constant; they are caught +by :func:`is_collapsed` and zeroed. +""" + + +def any_rank_has_non_finite(values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None) -> bool: + """Whether *any* rank's shard holds a NaN or an infinity. + + A local ``isfinite`` check in front of a collective is a deadlock: the rank + that finds the bad value raises and leaves, and every other rank blocks + forever in the reduction it never reaches. That is the same failure + :func:`relax.algorithms.advantages._agree_on_segmentation` exists to + prevent, and it is easy to reintroduce because the local check looks + self-contained. + + One MAX all-reduce of a 0/1 flag, so every rank gets the same answer and + the caller can raise on all of them together. Without a group the reduction + is skipped and the answer is simply the local one. + + The flag is built and reduced entirely on device. An earlier version read + it back with ``bool(...)`` to construct the tensor, which cost a + device-to-host sync *before* the collective and another host-to-device copy + to rebuild it -- and this runs on every training batch of GDPO's step 3. + The single ``.item()`` at the end is the only sync, and it is unavoidable + while the function returns a Python ``bool``. + + The flag is cast off ``bool`` on purpose, though not because bool is known + to break. MAX over a 0/1 flag gives the same answer in any of these dtypes, + and bool measurably survives both backends we can run: gloo in CI, and + NCCL 2.28.9 under torch 2.11 on two H100s, where bool, uint8, int32, int64 + and float32 all reduce correctly at both scalar and ``[1]`` shape. + + What we cannot run is the rest of the matrix. ``relax/utils/device.py`` + picks the backend from the accelerator, so an NPU job reduces this flag + over hccl and an XPU job over xccl, and neither vendor's dtype table is + something this repo can check. A dtype table is a lookup; a missing entry + throws rather than degrades, and it would throw on the first batch of + GDPO's step 3. A numeric flag costs nothing and takes the question off the + table -- that is the whole argument for the cast, and it is precaution + rather than a fix for any failure observed so far. + + float32 specifically is just what the pre-rewrite code sent; this was a + latency fix and had no business changing the wire format. ``int32`` would + be equally safe and would match every other flag in this repo + (megatron/actor.py, megatron/collective_utils.py, + megatron/conditional_branch_sync.py, all of which also send shape ``[1]`` + rather than a scalar) -- that alignment is worth doing, but as its own + change, not as a side effect of this one. + """ + bad = (~torch.isfinite(values)).any().to(dtype=torch.float32) + if process_group is not None: + dist.all_reduce(bad, op=dist.ReduceOp.MAX, group=process_group) + return bool(bad.item()) + + +def is_collapsed(values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None) -> bool: + """Whether every value is identical, i.e. the spread carries no signal. + + Tested by exact equality rather than by comparing the standard deviation + against a tolerance. A tolerance has to be relative to the magnitude (the + mean of N equal float32 values does not come back exactly equal to them, so + a collapsed group still shows std ~= 1e-8 times its magnitude), and any + relative tolerance large enough to catch that also erases real signal: with + ``std <= 1e-6 * max|x|``, the perfectly informative batch + ``[10000, 10000.005, 10000.010, 10000.015]`` is thrown away. Exact equality + has no such false positives, and near-equality is already damped by the + caller's epsilon (:data:`GDPO_EPS` on the GDPO whitening path, + :data:`STD_EPS` on the group path). + """ + if process_group is None: + if values.numel() == 0: + return True + return bool(values.min() == values.max()) + + # Every rank in the group has to reach the collective, including one whose + # shard came out empty — returning early there would hang the others. An + # empty shard contributes -inf to both halves, which is the identity for MAX + # and therefore leaves the reduction to the ranks that do have samples. + empty = values.numel() == 0 + neg_infinity = torch.tensor(float("-inf"), dtype=values.dtype, device=values.device) + bounds = torch.stack( + [neg_infinity if empty else -values.min(), neg_infinity if empty else values.max()], + ) + dist.all_reduce(bounds, op=dist.ReduceOp.MAX, group=process_group) + low, high = -bounds[0], bounds[1] + if not torch.isfinite(low): + return True # every rank was empty + return bool(low == high) + + +def collapsed_columns(values: torch.Tensor, dim: int) -> torch.Tensor: + """Per-column version of :func:`is_collapsed` for a ``[G, K]`` group. + + Returns a boolean tensor of shape ``[K]``: ``True`` where that reward + component took the same value across the whole group. + """ + return values.amax(dim=dim) == values.amin(dim=dim) + + +def distributed_mean_std( + values: torch.Tensor, *, process_group: dist.ProcessGroup | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + """Mean and unbiased std of ``values``, optionally across + ``process_group``. + + Each rank holds its own shard of the batch, so a local ``std()`` would give + every rank a different scale factor. Reducing across the group makes the + statistics describe the whole batch, which is what the framework already + does for ``--normalize-advantages`` (see + ``relax.utils.distributed_utils.distributed_masked_whiten``). + + Two passes, in float64. The one-pass form ``E[x^2] - E[x]^2`` subtracts two + nearly equal large numbers when the values sit far from zero, and the result + is dominated by rounding: for ``[1000.0, 1000.01, 1000.02, 1000.03]`` it + returns a variance of exactly 0 (true std 1.29e-2), and for the tighter + ``[10000.0, 10000.001, 10000.002, 10000.003]`` it returns std 4.6 instead of + 1.3e-3 — off by a factor of 3660. Scaling the first example up to 1e4 + without tightening it just reproduces the zero, not the inflation. + Neither is loud. The first silently zeroes every advantage in the batch; + the second silently rescales them. Reward magnitudes like these are + ordinary: the GDPO paper's own maths setup uses a length reward, and token + counts live in the thousands. + + The extra collective is two scalars, which is not worth optimising away. + """ + _log_group_once(process_group) + + # float64 throughout: the cancellation above is a precision problem, and + # doing the arithmetic in the caller's float32 reintroduces it even with the + # two-pass formula. + work = values.double() + count = torch.tensor(float(work.numel()), dtype=torch.float64, device=work.device) + total = work.sum() + + if process_group is not None: + first = torch.stack([count, total]) + dist.all_reduce(first, op=dist.ReduceOp.SUM, group=process_group) + count, total = first[0], first[1] + + if count == 0: + zero = torch.zeros((), dtype=values.dtype, device=values.device) + return zero, zero + + mean = total / count + # Centred before squaring, so no large offset survives into the sum. + sum_sq_dev = (work - mean).pow(2).sum() + if process_group is not None: + dist.all_reduce(sum_sq_dev, op=dist.ReduceOp.SUM, group=process_group) + + # Bessel-corrected, matching torch.std()'s default so the single-reward GDPO + # scale factor stays derivable from the GRPO group statistics. + variance = sum_sq_dev / torch.clamp(count - 1, min=1.0) + # variance cannot be negative now that it is a sum of squares; the clamp is + # only guarding the exactly-zero case against a -0.0. + std = torch.sqrt(torch.clamp(variance, min=0.0)) + return mean.to(values.dtype), std.to(values.dtype) diff --git a/relax/algorithms/policy.py b/relax/algorithms/policy.py new file mode 100644 index 000000000..dc2974123 --- /dev/null +++ b/relax/algorithms/policy.py @@ -0,0 +1,76 @@ +# 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, GDPO, 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..709d69dc0 --- /dev/null +++ b/relax/algorithms/rewards.py @@ -0,0 +1,678 @@ +# 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 — even for multi-reward algorithms, which collapse their components to a +single scalar here. +""" + +import math +from numbers import Real +from typing import Any, Callable + +import torch + +from relax.algorithms.numerics import GDPO_EPS, STD_EPS, collapsed_columns +from relax.algorithms.spec import get_algorithm +from relax.utils.logging_utils import get_logger +from relax.utils.training.ppo_utils import compute_rloo_leave_one_out_rewards + + +logger = get_logger(__name__) + +GROUP_EPS = STD_EPS + +_FLOAT32_MAX = float(torch.finfo(torch.float32).max) +"""Largest finite float32. Rewards are carried as float32 from here on, so a +value above this is an overflow waiting to happen rather than a large reward.""" + + +def group_positions(samples: list[Any], expected_size: int) -> dict[int, list[int]]: + """Map ``Sample.group_index`` to the positions it occupies in ``samples``. + + 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() + + +def resolve_gdpo_keys(args: Any) -> list[str]: + """The reward components GDPO normalises independently.""" + keys = list(getattr(args, "gdpo_reward_keys", None) or []) + if len(keys) < 2: + raise ValueError(f"--gdpo-reward-keys needs at least two reward keys, got {keys}.") + duplicates = {key for key in keys if keys.count(key) > 1} + if duplicates: + raise ValueError(f"--gdpo-reward-keys contains duplicates: {sorted(duplicates)}.") + return keys + + +def resolve_gdpo_weights(args: Any, keys: list[str]) -> list[float]: + """Per-component weights, defaulting to 1.0 each. + + The weights multiply the *normalised* advantages (arXiv 2601.05242, Eq. 7), + not the raw rewards. That is the point: after step 1 every component is on + the same scale, so a weight expresses relative importance rather than + accidentally encoding the component's units. + """ + weights = getattr(args, "gdpo_reward_weights", None) + if weights is None: + return [1.0] * len(keys) + if len(weights) != len(keys): + raise ValueError(f"--gdpo-reward-weights has {len(weights)} entries but --gdpo-reward-keys has {len(keys)}.") + resolved = [float(w) for w in weights] + for key, weight in zip(keys, resolved, strict=True): + # argparse happily parses "nan" and "inf" for a float option. Unchecked, + # the weighted sum turns non-finite and whiten_scalar refuses the batch. + # It used to read a non-finite std as a collapse and return zeros + # silently; catching the weight here still gives the better message. + if not math.isfinite(weight): + raise ValueError(f"--gdpo-reward-weights for {key!r} is {weight}, which is not finite.") + if all(weight == 0.0 for weight in resolved): + # Every component gets multiplied by zero, so the combined advantage is + # identically zero: the run trains on no signal and still exits cleanly. + # The same shape of failure as the non-finite case above, one line later. + raise ValueError(f"--gdpo-reward-weights are all zero ({resolved}); the combined advantage would be 0.") + return resolved + + +def extract_reward_components(samples: list[Any], keys: list[str]) -> torch.Tensor: + """Build the ``[B, K]`` component matrix, rejecting malformed rewards. + + Contract violations raise instead of defaulting to 0.0. A silently zeroed + component is indistinguishable from a genuinely collapsed one, so falling + back would hide a broken reward function behind plausible-looking training. + """ + rows: list[list[float]] = [] + for position, sample in enumerate(samples): + values = sample.get_reward_components(keys) + row: list[float] = [] + for key, value in zip(keys, values, strict=True): + # `numbers.Real` rather than `(int, float)`: reward functions routinely + # return numpy scalars, and only np.float64 happens to subclass float — + # np.float32/np.int64 would be rejected while np.float64 sailed through. + # bool and np.bool_ stay rejected (bool subclasses int; np.bool_ is not + # a Real), because a boolean reward is almost always a mistake. + if isinstance(value, bool) or not isinstance(value, Real): + raise TypeError( + f"Reward {key!r} of sample {position} must be a real number, " + f"got {value!r} ({type(value).__name__})." + ) + numeric = float(value) + if not math.isfinite(numeric): + raise ValueError(f"Reward {key!r} of sample {position} is {numeric}, which is not finite.") + # Finite in float64 is not enough: the tensor below is float32, whose + # largest value is ~3.4e38. A reward of 1e300 passes `isfinite` here, + # becomes `inf` on cast, and then reads as a non-finite std one stage + # later -- where the batch is zeroed and the run trains on no signal + # without ever failing. Catching the overflow at the boundary keeps + # the diagnosis at the reward function that produced it. + if abs(numeric) > _FLOAT32_MAX: + raise ValueError( + f"Reward {key!r} of sample {position} is {numeric!r}, which overflows float32 " + f"(max {_FLOAT32_MAX:.6g}). Rescale the reward; casting it would silently produce inf." + ) + row.append(numeric) + rows.append(row) + + # float64, not float32. The rewards arrive as Python floats and the + # standardisation that follows is ill-conditioned exactly when two + # components cancel: `C - r` loses the low bits on a float32 cast, and + # dividing what is left by a near-zero std turns those lost bits into a + # full-scale advantage. Keeping the width here is the only thing that keeps + # that residue below the signal at all: in float32 it is larger, and step 3 + # returns advantages of order 1 for a group that carries none. + # The float32 range check above still stands: the combined reward is cast + # back down further along, so a value that cannot survive that cast is + # still a bug worth reporting at the reward function that produced it. + components = torch.tensor(rows, dtype=torch.float64) + # Belt and braces: the per-value check above is exact, but it only sees what + # `float()` returned. Anything that slips past it must not reach the + # normaliser, where non-finite input is indistinguishable from a collapse. + if not torch.isfinite(components).all(): + bad = (~torch.isfinite(components)).nonzero()[0].tolist() + raise ValueError(f"Reward {keys[bad[1]]!r} of sample {bad[0]} is not finite as a tensor.") + return components + + +def standardize_group_components(group: torch.Tensor) -> torch.Tensor: + """GDPO step 1 on one ``[G, K]`` prompt group (arXiv 2601.05242, Eq. 4). + + Each column is standardised on its own. A column that is flat across the + group contributes exactly zero rather than ``0 / (0 + eps)`` noise, which + is what lets the remaining columns keep their signal. + + The arithmetic runs in float64, for the same reason + :func:`relax.algorithms.numerics.distributed_mean_std` does: the means and + stds are the part where cancellation bites, and widening them is cheap. + + Standardising is ill-conditioned when a column barely varies: it divides + by a std that is near zero, so the *relative* error in the result grows + like ``max|x| / std``. That is not a corner case here -- it is precisely + what two columns summing to a constant look like, and it is why the columns + arrive here in float64: in float32 the residue of that cancellation was + larger than the signal it was left over from. See :func:`combine_group` for + the part of this that is still unsolved. + """ + work = group.double() + centered = work - work.mean(dim=0, keepdim=True) + std = work.std(dim=0) + collapsed = collapsed_columns(group, dim=0) + scaled = centered / (std + GDPO_EPS) + scaled = torch.where(collapsed.unsqueeze(0), torch.zeros_like(scaled), scaled) + return scaled.to(group.dtype) + + +def combine_group(group: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + """GDPO steps 1 and 2 for one ``[G, K]`` prompt group: ``sum_k w_k z_ik``. + + Eq. 7 and nothing else. Two things used to sit at the end of this function + and both are gone; the reasons are worth keeping because both sounded good. + + **The noise floor** returned zeros when the result fell below + ``8 * sum(|w_k| * noise_k)``. It is not in Eq. 7, and the tests written to + document it proved it destroyed real signal: two components at + ``base = 4.05e13``, every column measuring under 1% noise, gave a floor of + 0.148 against a true signal of 0.033. It grew as a sum over components + while any per-column noise measure is a max, so sixteen clean columns + reached 0.82. + + **Its replacement, a relative "is this rounding?" ratio**, lasted one round + longer and was worse, because it was wrong in kind rather than in + calibration. ``magnitude`` scales with the *difference* between the weights + while ``sum_k |w_k| max|z_k|`` scales with their *magnitude*. Measured on + real inputs it was inverted in both directions at once -- it discarded a + group whose final advantage was 0.43 and kept one whose 1.08 was pure + rounding. + + That is not fixable by moving the threshold. For ``G >= 3`` the centred + subspace is at least two-dimensional, so ``z_2 = -z_1 + delta * u`` with + ``u`` orthogonal to ``z_1`` and ``delta`` arbitrarily small is a genuine + signal with an arbitrarily small ratio. No universal threshold separates + "the components nearly cancel" from "the arithmetic nearly cancelled". + + ``G = 2`` does not settle it either way: the ratio reduces to + ``|w1 - w2| / (|w1| + |w2|)``, independent of the data, *only* when the two + standardised columns come out as exact opposites. Columns that move + together give 1 for same-signed weights, which is fully data-dependent. The + ``G >= 3`` construction above is what actually settles it. + + **What is left unsolved, stated plainly rather than argued away.** Two + components summing to a constant do not cancel exactly, and the remainder + grows with the magnitude they are centred on -- roughly ``ulp(C) / spread`` + for the Eq. 7 output. Nothing in this file detects it, and the two + mechanisms tried so far could not: a conditioning bound cannot prove a + small result is *not* signal, and screening on the post-whitening magnitude + needs the whole batch, which a per-group function does not have. + + What reaches the optimizer, measured (``[C - d, d]`` for + ``d in (0.1, 0.2, 0.3, 0.7)``, equal weights): + + ============ ========================= =================================== + ``C`` its own whitening unit sharing one with 8 healthy groups + ============ ========================= =================================== + ``1e8`` 2.3e-4 2.7e-8 + ``1e9`` 2.6e-3 3.1e-7 + ``1e11`` 2.0e-1 3.3e-5 + ``1e13`` 1.2 5.2e-3 + ============ ========================= =================================== + + The second column is not a bound. Sharing a whitening unit with healthy + groups divides the remainder by *their* standard deviation -- a constant + factor of a few hundred here. It does not slow the growth with ``C``, and + at ``C = 1e13`` a mixed unit still delivers 5e-3. + + "Whitening unit" is also not "the rollout". :func:`~relax.algorithms. + advantages._whiten_by_segment` whitens each *training batch* separately, so + a degenerate group only gets that division if healthy groups land in the + same training batch. Isolate it into its own segment and it returns to the + first column -- 2.6e-3 at ``C = 1e9``. Nothing arranges that mixing; it is + a property of how the rollout happened to be split. + """ + standardized = standardize_group_components(group) + return (standardized.double() * weights.double()).sum(dim=1) + + +def normalize_gdpo_decoupled(args: Any, samples: list[Any], raw_rewards: list[float]) -> list[float]: + """GDPO steps 1 and 2 (arXiv 2601.05242, Eq. 4 and Eq. 7). + + Step 1 standardises each reward component within its prompt group; step 2 + combines them with the configured weights. The result is one scalar per + sample, so it travels through the existing ``rewards`` column and needs no + TransferQueue schema change. Step 3 (batch whitening) runs later, in + :func:`relax.algorithms.advantages.advantage_gdpo`. + + What standardising per component actually buys, stated carefully because + it is easy to overclaim in two opposite directions (this docstring has + managed both): + + The combined advantage is ``sum_k w_k * z_k`` with each ``z_k`` at unit + variance, so its variance is ``sum_k w_k^2 + 2 sum_{i +1``) amplify, contradicting ones (``rho -> -1``) attenuate and + at the limit cancel. + + Both claims this file got wrong earlier are special cases of that one + formula. 'Equal sums are rescued by GDPO' is ``rho = -1`` -- they cancel, + GDPO included. 'Two varying components mean a stronger signal' is + ``rho = +1`` -- true there, false at ``rho = -0.8``, where the combination + measures 0.63x a single component. + + Separately and unconditionally: GRPO's direction is dominated by whichever + component has the largest spread, because it standardises the raw sum. A + correctness reward in ``{0, 1}`` added to a length reward in the hundreds + is decided almost entirely by length under GRPO, and half by each here. + """ + keys = resolve_gdpo_keys(args) + weights = resolve_gdpo_weights(args, keys) + + components = extract_reward_components(samples, keys) + positions_by_group = group_positions(samples, args.n_samples_per_prompt) + + # float64, matching the components. The weights used to be cast to float32 + # here, one stage before the arithmetic needed it, and that quantisation was + # observable: [16777216, 16777217] are two distinct configured weights that + # become the same float32, so a pair of anti-correlated components that + # should combine to about +-1 combined to exactly 0 instead. The cast to the + # transport dtype belongs at the end of the pipeline, not at the start. + weight_tensor = torch.tensor(weights, dtype=torch.float64) + + # The two checks below are still about float32, and deliberately: the + # combined reward is carried to the trainer in float32 (`dict_to_tensordict`), + # so a weight vector that cannot survive that cast produces a run that trains + # on nothing. Asking here rather than after the multiplication is what lets + # the message name the weights instead of the arithmetic. + if not torch.isfinite(weight_tensor.float()).all(): + # `resolve_gdpo_weights` checked `math.isfinite` on the Python floats. + # That is a different question: 1e300 is finite in float64 and `inf` + # once it reaches the transport dtype. + raise ValueError( + f"--gdpo-reward-weights {weights} contains a value that overflows float32; " + f"the largest representable is {_FLOAT32_MAX:.6g}." + ) + if float(weight_tensor.float().abs().sum()) == 0.0: + # Same gap in the other direction: 1e-50 is a nonzero Python float that + # flushes to zero in float32, so a weight vector that passed the + # "not all zero" check can still be all zeros by the time it matters. + raise ValueError( + f"--gdpo-reward-weights {weights} are all zero once cast to float32; " + "the combined advantage would be identically 0." + ) + + combined = torch.zeros(len(samples), dtype=torch.float64) + silent_groups = 0 + for positions in positions_by_group.values(): + group_combined = combine_group(components[positions], weight_tensor) + combined[positions] = group_combined + # The same predicate `group_carries_reward_signal` applies, on the same + # float32 view. It used to be `.any()` on the float64 values, which is a + # different question -- the warning and the sampler could disagree about + # the same group, and the one the operator sees is the warning. + as_transported = group_combined.float() + if bool(as_transported.amin() == as_transported.amax()): + silent_groups += 1 + + if silent_groups == len(positions_by_group): + # Every group came out at exactly zero, so this batch produces no + # gradient at all. Usually the reward function is constant for these + # prompts (e.g. a format reward when nothing in the prompt asks for a + # format); the other way in is components that cancel. Worth one line, + # because the symptom downstream is simply "loss does not move". + logger.warning( + "GDPO: all %d groups of this batch combined to a value that does not vary within the " + "group once cast to float32 (keys=%s); the batch contributes no gradient. Either these " + "rewards do not vary across rollouts of the same prompt, or they cancel under the " + "configured weights (weights=%s).", + silent_groups, + keys, + weights, + ) + + if not torch.isfinite(combined).all(): + # Every individual reward fit in float32 (checked in + # `extract_reward_components`), but the arithmetic between them need + # not. The mean itself no longer overflows -- the columns and the + # standardisation are float64 now, so [3e38, 2e38, 1e38, 0] averages + # fine, and the earlier version of this comment claiming otherwise is + # out of date. What can still reach infinity is the weighting: float64 + # weights near its own maximum, or a reward near float64's range rather + # than float32's. Left unchecked the `inf` reaches `whiten_scalar`, + # which now raises on it -- but at that point the message can only name + # the batch, not the reward that produced it. + raise ValueError( + "GDPO produced a non-finite combined advantage from finite inputs; the " + f"arithmetic overflowed. Rescale the rewards (keys={keys}) or their weights." + ) + return combined.tolist() + + +def group_carries_reward_signal(args: Any, samples: list[Any]) -> bool: + """Whether one prompt group still carries signal for *this* algorithm. + + Consumers upstream of the advantage stage -- the dynamic-sampling filter, + the zero-std metrics -- ask "did this group vary?" to decide whether it is + worth keeping or worth reporting. They have always answered it with the + single scalar ``--reward-key`` selects, which is the right question only + for an algorithm that consumes that scalar. + + For a multi-reward algorithm the honest test is not "did any component + vary" but "is the combined advantage this algorithm will actually compute + vary in float32". Those differ: a zero weight mutes a varying component, and two + components whose standardised values are exact opposites cancel. Asking + the cheaper question keeps groups that then contribute no gradient, and + under-reports the zero-std metrics. Steps 1 and 2 are per-group and cheap, + so this runs them. + + The single-reward test is deliberately performed **in float32**, on a + tensor, rather than on the Python floats. That is the precision the reward + actually reaches training in (:func:`_group_normalize` casts to float32, + and so does the GDPO path), so it is the precision that decides whether a + group will still carry signal by the time it matters. + + Comparing the float64 values instead is not a harmless tightening: it flips + the verdict for any group whose spread survives in float64 but collapses on + cast -- ``[0.1 + 0.2, 0.3, ...]`` is the everyday example, and a group of + NaNs is the pathological one (``nan != nan`` reads as signal). Both would + be kept and then contribute a zero gradient. Judging on the tensor keeps + this identical to the ``std > 0`` test it replaces: ``min == max`` and + ``std == 0`` agree on every float32 input. + """ + if not samples: + return False + + spec = get_algorithm(args.advantage_estimator) + if not spec.uses_reward_components: + rewards = torch.tensor([sample.get_reward_value(args) for sample in samples], dtype=torch.float32) + if not torch.isfinite(rewards).all(): + # A group containing NaN or inf reads as "varying" under any + # inequality test (`nan != nan` is True), which would forward a + # broken reward into training. `std > 0` answered False here + # because the std is itself NaN, so False preserves that. Raising + # would arguably be better -- a non-finite reward is a bug in the + # reward function, not a property of the group -- but that is a + # behaviour change this refactor is not the place for. + return False + return bool(rewards.amin() != rewards.amax()) + + keys = resolve_gdpo_keys(args) + # float64, the same dtype `normalize_gdpo_decoupled` weights with. A float32 + # weight tensor here would answer this question for a *different* weight + # vector than the one training uses, whenever two configured weights are + # closer together than float32 can represent. + weights = torch.tensor(resolve_gdpo_weights(args, keys), dtype=torch.float64) + components = extract_reward_components(samples, keys) + # `combine_group`, not steps 1 and 2 open-coded: this has to answer the same + # question `normalize_gdpo_decoupled` will. Open-coding it is how the two + # came apart before. + # + # The same *form* the single-reward branch above uses, on the values the + # trainer will receive: `min != max` in float32. Not the same question, + # though -- that branch reads the raw reward, this one reads the output of + # Eq. 4 and Eq. 7, and the two disagree on a group whose components sum to + # a constant (the scalar is flat, the combination is rounding). + # + # `min != max` rather than `.any()` because it matches that branch, and + # here the two coincide anyway: Eq. 4 centres every column, so the combined + # values sum to zero within the group and "all equal" forces "all zero". + # That is a borrowed invariant, not a property of this function -- + # `test_eq7_centres_every_group` pins it, because without it a group could + # be constant and nonzero, and such a group is *not* whitened away by step + # 3 (step 3 works per training batch, so it would come out as a constant + # nonzero advantage). + # + # No tolerance of any kind: a threshold here decides whether a prompt group + # enters training, and no threshold on this quantity can tell a genuine + # near-cancellation from a rounding one (see `combine_group`). + # + # The consequence is deliberate and it has a cost. A group whose components + # sum to a constant survives this test on its rounding remainder and is + # trained on, and it wastes the group. How much reaches the optimizer grows + # with the magnitude the components are centred on -- see the table in + # `combine_group`; it is 3.1e-7 for C = 1e9 in a mixed training batch but + # 5.2e-3 at C = 1e13, and a degenerate group that lands in a batch of its + # own gets no division at all. That is the price of not guessing. + combined = combine_group(components, weights).float() + return bool(combined.amin() != combined.amax()) + + +def observed_reward_signal(args: Any, samples: list[Any]) -> bool | None: + """:func:`group_carries_reward_signal` for callers that only report. + + Returns ``None`` -- "cannot tell" -- where the strict version raises. + + This is not the strict check with its errors swallowed; it is a different + question, and the difference is not cosmetic. The strict version is asked + by the dynamic-sampling filter, which decides whether a group enters + *training*, so a reward it cannot read is a bug that should stop the run. + The zero-std metrics are asked by the logger, and two things follow: + + * A metric must not decide whether training proceeds. Making observation + the first enforcer of a contract means a broken reward function is + reported at the log line rather than at the stage that consumes it, and + it takes the rollout down on the way. + * The contract is not even in force everywhere the metrics run. Eval may + use a different reward model (``EvalConfig.rm_type``), so an eval reward + legitimately need not carry the ``--gdpo-reward-keys`` that training + needs -- nothing in eval consumes them. The strict question has no + correct answer there; ``None`` is the honest one. + + During training the same violation is still raised, by + :func:`normalize_gdpo_decoupled`, which is the stage that actually reads + the components. Nothing is hidden -- but it is logged here too, because a + group the metrics could not read is worth knowing about either way. + """ + try: + return group_carries_reward_signal(args, samples) + except (TypeError, ValueError, KeyError, IndexError) as exc: + # KeyError/IndexError are not hypothetical. `get_reward_components` + # raises ValueError for a missing key, but the single-reward branch + # goes through `Sample.get_reward_value`, which is a bare subscript + # (`relax/utils/types.py:171`): a `--reward-key` absent from the reward + # dict raises KeyError straight through this handler and takes the + # metrics path down -- the exact failure this function exists to + # prevent. Narrowing `get_reward_value` to ValueError would be the + # tidier fix, but that accessor is shared with the filters and the + # rollout metrics, so changing what it raises is a contract change for + # callers that are not in scope here. + logger.warning( + "zero-std metrics: skipping a group whose reward could not be read (%s). " + "For a training rollout the reward stage will raise on this; for eval it may just " + "mean the eval reward model returns a different schema, which is fine.", + exc, + ) + return None + + +def metrics_group_verdict(args: Any, samples: list[Any]) -> bool | None: + """Is this prompt group flat, for zero-std reporting? ``None`` if + unanswerable. + + ``True`` flat, ``False`` varying, ``None`` neither -- nothing in the group + was scored, or the reward could not be read. + + This lives here rather than in the two rollout modules because there are + *two* copies of ``_compute_zero_std_metrics``, one in + ``relax/agentic/rollout.py`` and one in ``relax/distributed/ray/rollout.py``, + and they have already drifted apart once: the agentic one dropped unscored + samples and the distributed one did not, which turned a `reward=None` into + a ``TypeError`` on the rollout's way out. Only the agentic one is reachable + from a CPU test -- the other pulls in ``sglang`` at import -- so a shared + helper is the only version of this logic that can be tested at all. + + ``reward=None`` is a real state, not a defensive check: under ``--group-rm`` + the group reward is assigned in one shot that is skipped entirely when the + rollout aborts, which is why the "reward is not None" assert in + ``sglang_rollout.py`` exempts ``group_rm`` in the first place. + + Callers still differ on what to do with ``None``, and that difference is + deliberate -- each preserves what its own metric reported before. See the + call sites. + """ + rewarded = [sample for sample in samples if sample.reward is not None] + if not rewarded: + return None + signal = observed_reward_signal(args, rewarded) + return None if signal is None else not signal + + +def zero_std_group_label(args: Any, samples: list[Any]) -> str | None: + """The ``zero_std/count_