【Task.27】feat(algorithms): add GDPO multi-reward advantage estimator

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' },

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
 

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/<algo>/`: 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/`.

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 |