【Task 40】Feat(p3o): add ESS-adaptive policy optimization

docs/en/examples/algorithms.md

@@ -277,11 +277,46 @@ SAPO_ARGS=(
 
 ---
 
+## P3O
+
+P3O corrects rollout-policy mismatch with selected-token behavior
+log-probabilities. It computes an effective sample size
+(ESS) from the importance ratios and uses the detached ESS value as a one-sided
+adaptive cap for the policy update.
+
+P3O is mutually exclusive with `--use-opd`: combining its objective with an OPD
+teacher loss or OPD advantage replacement would create an unvalidated hybrid.
+
+### Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--advantage-estimator p3o` | — | Enable P3O |
+| `--use-rollout-logprobs` | required | Use rollout behavior log-probabilities for importance ratios |
+| `--calculate-per-token-loss` | required | Preserve P3O's token-sum normalization |
+| `--p3o-ess-scope` | `micro-batch` | Compute the adaptive cap per micro-batch; `step` is available for replay validation |
+| `--p3o-kl-mode` | `proxy` | Behavior-KL approximation: `proxy` or `proxy_safe` |
+| `--clip-low`, `--clip-high` | `0.2` | P3O clip-fraction monitoring margins |
+
+The full-vocabulary `exact` KL calculation is a pure verification helper, not a
+production CLI mode: rollout records contain selected-token behavior
+log-probabilities rather than full behavior logits.
+
+### Recipes
+
+The A100×4 recipes pair P3O and GRPO with identical on-policy, periodic-sync,
+and temperature-mismatch scenarios. Start with
+`examples/algorithms/p3o/README.md`; use
+`examples/algorithms/p3o/run_p3o_smoke.sh` for a one-rollout smoke check.
+
+---
+
 ## Algorithm Comparison
 
 | Algorithm | Advantage Computation | Policy Loss | KL Constraint |
 |-----------|----------------------|-------------|---------------|
 | **PPO** | Critic values + GAE | PPO-Clip (hard clip) | Disabled in the current synchronous topology |
+| **P3O** | Rollout behavior log-probabilities + ESS | Detached one-sided adaptive cap | Sampled-token behavior-KL proxy |
 | **GRPO** | Group-relative reward | PPO-Clip (hard clip) | Optional KL loss |
 | **REINFORCE++** | Token KL-to-go return + global token normalization | PPO-Clip (hard clip) | k1 KL in shaped reward |
 | **REINFORCE++-baseline** | Inclusive group mean + global token normalization | PPO-Clip (hard clip) | Separate k2 KL loss |
@@ -295,5 +330,6 @@ SAPO_ARGS=(
 - [PPO Training](../guide/ppo-training.md)
 - [REINFORCE++ Training](../guide/reinforce-plus-plus.md)
 - [Quick Start](../guide/quick-start.md)
+- `examples/algorithms/p3o/README.md`
 - [On-Policy Distillation](./on-policy-distillation.md)
 - [Generative Reward Model](./generative-reward-model.md)

docs/zh/examples/algorithms.md

@@ -274,11 +274,43 @@ SAPO_ARGS=(
 
 ---
 
+## P3O
+
+P3O 使用 rollout 中记录的已选 token 行为策略 log-probability 校正
+rollout-policy mismatch。它由重要性比率计算有效样本量
+（ESS），并将 detach 后的 ESS 用作单侧自适应策略更新上限。
+
+P3O 与 `--use-opd` 互斥：把 P3O 目标与 OPD teacher loss 或 OPD advantage
+replacement 组合会形成未经验证的混合目标。
+
+### 关键参数
+
+| 参数 | 默认值 | 说明 |
+|------|--------|------|
+| `--advantage-estimator p3o` | — | 启用 P3O |
+| `--use-rollout-logprobs` | 必需 | 使用 rollout 行为策略 log-probability 计算重要性比率 |
+| `--calculate-per-token-loss` | 必需 | 保持 P3O 的 token-sum 归一化 |
+| `--p3o-ess-scope` | `micro-batch` | 每个 micro-batch 计算自适应上限；`step` 用于 replay 验证 |
+| `--p3o-kl-mode` | `proxy` | 行为 KL 近似：`proxy` 或 `proxy_safe` |
+| `--clip-low`、`--clip-high` | `0.2` | P3O clip-fraction 监控边距 |
+
+全词表 `exact` KL 计算是纯验证辅助函数，并非生产命令行模式：rollout 记录的是
+已选 token 的行为策略 log-probability，而非完整行为 logits。
+
+### 配置示例
+
+A100×4 配置会以相同的 on-policy、周期同步和温度失配场景成对比较 P3O 与 GRPO。
+请从 `examples/algorithms/p3o/README_zh.md` 开始；一轮 rollout 的冒烟检查使用
+`examples/algorithms/p3o/run_p3o_smoke.sh`。
+
+---
+
 ## 算法对比
 
 | 算法 | Advantage 计算 | 策略损失 | KL 约束方式 |
 |------|---------------|---------|-----------|
 | **PPO** | Critic value + GAE | PPO-Clip（硬裁剪） | 当前同步拓扑中禁用 |
+| **P3O** | Rollout 行为策略 log-probability + ESS | detach 的单侧自适应上限 | 已选 token 行为 KL proxy |
 | **GRPO** | 组相对奖励 | PPO-Clip（硬裁剪） | 可选 KL loss |
 | **REINFORCE++** | Token KL-to-go return + 全局 token 归一化 | PPO-Clip（硬裁剪） | shaped reward 中的 k1 KL |
 | **REINFORCE++-baseline** | Inclusive group mean + 全局 token 归一化 | PPO-Clip（硬裁剪） | 独立 k2 KL loss |
@@ -292,5 +324,6 @@ SAPO_ARGS=(
 - [PPO 训练](../guide/ppo-training.md)
 - [REINFORCE++ 训练](../guide/reinforce-plus-plus.md)
 - [快速开始](../guide/quick-start.md)
+- `examples/algorithms/p3o/README_zh.md`
 - [在线策略蒸馏](./on-policy-distillation.md)
 - [生成式奖励模型](./generative-reward-model.md)

examples/algorithms/README.md

@@ -4,20 +4,21 @@
 
 ## 概述
 
-Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator` 参数选择。GRPO、RLOO、CISPO、GSPO 与 SAPO 共享 Actor/Rollout 服务拓扑；RLOO 仅支持同步固定批量，PPO 还需要 Critic 与 Advantages 服务，应分别遵守专用配置。
+Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator` 参数选择。GRPO、P3O、RLOO、CISPO、GSPO 与 SAPO 共享 Actor/Rollout 服务拓扑；RLOO 仅支持同步固定批量，PPO 还需要 Critic 与 Advantages 服务，应分别遵守专用配置。
 
 ## 支持的算法
 
-| 算法                     | 启用参数                                             | 推荐场景                    |
-| ------------------------ | ---------------------------------------------------- | --------------------------- |
-| **PPO**                  | `--advantage-estimator ppo`                          | Actor-Critic、token 级 GAE  |
-| **GRPO**                 | `--advantage-estimator grpo`                         | 默认、大多数场景            |
-| **REINFORCE++**          | `--advantage-estimator reinforce_plus_plus`          | token KL-to-go 与全局归一化 |
-| **REINFORCE++-baseline** | `--advantage-estimator reinforce_plus_plus_baseline` | group baseline 与独立 k2 KL |
-| **RLOO**                 | `--advantage-estimator rloo`                         | 无偏基线、非裁剪 REINFORCE  |
-| **CISPO**                | `--advantage-estimator cispo`                        | 保留梯度方向、需要更高精度  |
-| **GSPO**                 | `--advantage-estimator gspo`                         | 序列级约束、稳定训练        |
-| **SAPO**                 | `--advantage-estimator sapo`                         | 平滑优化、soft 信任域       |
+| 算法                     | 启用参数                                             | 推荐场景                         |
+| ------------------------ | ---------------------------------------------------- | -------------------------------- |
+| **PPO**                  | `--advantage-estimator ppo`                          | Actor-Critic、token 级 GAE       |
+| **P3O**                  | `--advantage-estimator p3o`                          | 行为策略失配校正、ESS 自适应上限 |
+| **GRPO**                 | `--advantage-estimator grpo`                         | 默认、大多数场景                 |
+| **RLOO**                 | `--advantage-estimator rloo`                         | 无偏基线、非裁剪 REINFORCE       |
+| **REINFORCE++**          | `--advantage-estimator reinforce_plus_plus`          | token KL-to-go 与全局归一化      |
+| **REINFORCE++-baseline** | `--advantage-estimator reinforce_plus_plus_baseline` | group baseline 与独立 k2 KL      |
+| **CISPO**                | `--advantage-estimator cispo`                        | 保留梯度方向、需要更高精度       |
+| **GSPO**                 | `--advantage-estimator gspo`                         | 序列级约束、稳定训练             |
+| **SAPO**                 | `--advantage-estimator sapo`                         | 平滑优化、soft 信任域            |
 
 ## 选择建议
 
@@ -35,6 +36,13 @@ Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator
 - 对超出信任域的 token 直接置零梯度
 - 适合大多数强化学习场景
 
+### P3O（行为策略失配校正）
+
+- 使用 rollout 时记录的行为策略 log-probability 计算重要性比率，并以 ESS 自适应地限制策略更新
+- 生产模式只支持 `--p3o-kl-mode proxy` 与 `proxy_safe`；全词表 `exact` 仅供验证辅助函数使用，不是命令行模式
+- 与 `--use-opd` 互斥，且需要 `--use-rollout-logprobs` 与 `--calculate-per-token-loss`
+- A100×4 的成对 P3O/GRPO 场景与环境变量说明参见 [P3O 配置](p3o/README_zh.md)
+
 ### REINFORCE++ 两个变体
 
 - `reinforce_plus_plus` 使用 token 级 k1 KL reward shaping、KL-to-go return 和跨 DP rank 的有效 token 全局归一化

examples/algorithms/p3o/README.md

@@ -0,0 +1,200 @@
+# P3O A100×4 recipes
+
+These launchers compare P3O and GRPO under matched on-policy and controlled
+rollout-mismatch scenarios. They target four colocated GPUs and submit the
+training driver through Ray Jobs.
+
+中文版请参阅 [README_zh.md](README_zh.md)。
+
+## Required environment
+
+Set these paths before a non-dry run:
+
+```bash
+export P3O_MODEL_DIR=/path/to/model
+export P3O_TRAIN_DATA=/path/to/train.jsonl
+export P3O_EVAL_DATA=/path/to/eval.jsonl
+export P3O_OUTPUT_ROOT=/path/to/output
+export P3O_MEGATRON_DIR=/path/to/Megatron-LM
+export P3O_RAY_DASHBOARD=http://ray-dashboard-host:8265
+```
+
+`P3O_EVAL_DATA` is required in `formal` mode and is optional in `smoke` mode.
+The model, training data, and Megatron paths must exist before the Ray job is
+submitted. Each run records its resolved arguments, command, Git identity,
+logs, Ray status, exit code, and per-step rollout JSONL beneath
+`P3O_OUTPUT_ROOT`. Set `P3O_ROLLOUT_RESULT_DIR` only when an external evidence
+layout requires a different raw-rollout destination; the resolved path is
+recorded in `run_identity.env`.
+
+The Ray job runtime explicitly disables inherited HTTP proxies. SGLang checks
+engine health and registers workers through node-local IP addresses; allowing
+host proxy variables into Ray workers can leave healthy engines stuck behind
+the proxy instead of completing the startup barrier.
+
+Formal mode defaults to DeepScaleR's `problem`/`answer` fields. Smoke mode
+defaults to the commonly used `question`/`answer` schema. Set `P3O_INPUT_KEY`
+and `P3O_LABEL_KEY` explicitly when the selected asset uses another schema;
+both resolved keys are recorded in `run_identity.env`.
+
+Formal mode also defaults to the `deepscaler` rule-based verifier, which reads
+Qwen-Thinking's `</think>` suffix and a final `\\boxed{...}` answer. Smoke mode
+retains the `mopd` default for legacy GSM8K-style assets. Set `P3O_RM_TYPE`
+explicitly when a smoke uses DeepScaleR or another reward contract; the
+resolved reward type is recorded in `run_identity.env`.
+
+Formal evaluation defaults to the `deepscaler` dataset name, 16 samples per
+prompt, a 4096-token response cap, temperature 1.0, and top-p 0.95. Bounded
+resource studies may set `P3O_EVAL_NAME`, `P3O_EVAL_N_SAMPLES`,
+`P3O_EVAL_MAX_RESPONSE_LEN`, `P3O_EVAL_TEMPERATURE`, and `P3O_EVAL_TOP_P`.
+These values affect evaluation only and are recorded in `run_identity.env`;
+paired algorithms must use identical values.
+
+The default `P3O_ROLLOUT_SHUFFLE=1` retains ordinary training behavior. Set it
+to `0` only with a pre-materialized fixed prompt schedule for paired evidence;
+the setting is recorded so a shuffled run cannot be mistaken for the fixed
+comparison.
+
+Set `P3O_DETERMINISTIC_INFERENCE=1` for paired experiments that require common
+per-sample sampling seeds across P3O and GRPO. The resolved flag is recorded in
+run identity. This controls sampling randomness only; after the first update,
+different policy weights can and should produce different responses for the
+same seed.
+
+Formal mode sources `scripts/models/qwen3-4B.sh` and targets
+Qwen3-4B-Thinking-2507. Smoke mode sources `scripts/models/qwen3-0.6B.sh`.
+Set `P3O_MODEL_CONFIG` only when deliberately validating another compatible
+model configuration; the resolved path is recorded in `run_identity.env`.
+The formal launcher overrides the generic 4B script's RoPE base to `5000000`,
+matching this checkpoint's `config.json`; smoke remains at `1000000`. A
+deliberate compatible override can use `P3O_MODEL_ROTARY_BASE`, and its value is
+also recorded in run identity.
+
+## Active P3O contract
+
+The formal P3O path uses `--p3o-ess-scope micro-batch`,
+`--p3o-kl-mode proxy_safe`, and monitoring margins
+`--clip-low/--clip-high 0.2`. `proxy_safe` has the same forward value as the
+FeynRL-compatible sampled-token proxy and corrects only the extreme negative
+log-ratio gradient. `exact` is available only through the pure full-vocabulary
+verification helper, not as a CLI mode, because rollout data stores
+selected-token log-probabilities rather than behavior logits.
+
+P3O owns a dedicated policy-loss dispatch and is mutually exclusive with
+`--use-opd`; an OPD teacher loss, OPD advantage replacement, or OPD-only
+reward would define an unvalidated hybrid objective. The reward/verifier name
+`P3O_RM_TYPE=mopd` is unrelated to the `--use-opd` training feature and remains
+valid for compatible datasets.
+
+With `--context-parallel-size > 1`, P3O automatically runs every THD
+TransformerLayer, the final normalization, and the LM head on the reconstructed
+full sequence before slicing the result back to each CP rank. This strict path
+applies to both `micro-batch` and `step` ESS scopes and makes their forward and
+backward kernels see the same token order and shape as CP1. It duplicates
+full-sequence compute and activations on every CP rank, so peak memory and
+compute are higher than native context parallelism; whole-layer activation
+recomputation is recommended for long contexts. The supported contract is
+currently standard zig-zag THD with tensor parallel size 1. If the
+full-sequence path runs out of memory, P3O aborts and refuses CP>1 instead of
+silently falling back to the non-equivalent native CP kernel order.
+
+Strict partition invariance also requires Megatron `--deterministic-mode` and
+`--batch-invariant-mode`. The recipes enable both for matched P3O/GRPO runs and
+set `NCCL_ALGO=Ring`, `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`, and
+`CUBLAS_WORKSPACE_CONFIG=:4096:8` in the Ray runtime before model construction.
+P3O fails closed if either mode is absent. It also disables fused
+weight-gradient accumulation in this mode because the bundled
+batch-invariant TE GEMM cannot honor its multi-micro-batch `main_grad`
+accumulation contract. This uses the stable TE/DDP accumulation path and may
+reduce throughput or increase transient gradient memory.
+
+When data parallel size is greater than one, strict P3O also gathers the final
+training batch in global DP-rank order on every replica. DP rank zero keeps the
+real loss masks; the other replicas execute the same micro-batch schedule with
+zero loss masks so all DP/CP collectives remain aligned. The existing summed
+DP×CP reduction then propagates one canonical CUDA backward result. This
+removes rank-local backward drift, but deliberately duplicates forward/backward
+compute and provides no DP training-speedup in strict mode.
+The strict path also uses Megatron's eager fused-cross-entropy helpers: wrappers
+decorated with `torch.compile` at import time are explicitly unwrapped because
+their Triton autotune launcher is not valid when multiple replicas execute the
+same canonical micro-batch concurrently.
+The current strict-DP contract requires one rollout mini per optimizer step and
+fails closed instead of silently reordering multiple mini boundaries.
+Within each CP group, strict mode reconstructs the CP1 token order at the
+embedding, TransformerLayer, final-norm, LM-head, and P3O-loss boundaries.
+Forward values are broadcast from CP rank zero; backward reduces token-shard
+gradients to that same canonical graph and zeros the other replica graphs.
+This makes dense and embedding parameter accumulation run once in CP1 order
+instead of averaging rank-local CUDA or sparse-embedding results.
+
+Formal defaults are G=16, global batch 64, micro-batch 1, rollout batch 4,
+response length 4096, and 30 optimizer steps (`--num-rollout 30`). The planned
+paired seeds are 42, 123, and 2026. Smoke remains G=4, global batch 16,
+response length 128, and one optimizer step.
+
+The following environment variables expose the aligned settings without
+changing scenario scripts:
+
+```bash
+export P3O_ESS_SCOPE=micro-batch  # or step for capability/replay validation
+export P3O_KL_MODE=proxy_safe     # proxy for golden parity
+export P3O_CLIP_LOW=0.2
+export P3O_CLIP_HIGH=0.2
+export P3O_SEED=42
+export P3O_RM_TYPE=deepscaler # required when smoke mode is paired with DeepScaleR
+```
+
+Ray workers inherit normal proxy settings by default. On clusters where an
+injected outbound proxy intercepts SGLang's node-local readiness probes, set
+`P3O_CLEAR_RUNTIME_PROXIES=1` to clear proxy variables inside the job runtime.
+This setting is opt-in and recorded in `run_identity.env` because it also
+disables proxy access for every worker in the job.
+
+If A100-40GB capacity prevents a 4B pilot, reduce pilot response length first
+while keeping micro-batch size 1 and record the deviation. Do not treat reduced
+smoke runs as formal evidence or silently reduce the three-seed comparison.
+For a response-preserving resource fallback, set `P3O_ACTIVATION_RECOMPUTE=1`
+to add whole-layer uniform activation recomputation and set
+`P3O_LOG_PROBS_CHUNK_SIZE` to a positive token count for chunked log-probability
+and entropy reductions. Both settings apply identically to P3O and GRPO and are
+recorded in run identity; the default `0`/`-1` leaves the original path intact.
+
+## Scenarios
+
+| Scenario                   | Update interval | Temperature override | Meaning                                                           |
+| -------------------------- | --------------: | -------------------: | ----------------------------------------------------------------- |
+| `on_policy`                |               1 |                  off | Synchronize every rollout with the normal sampling configuration. |
+| `periodic_sync_interval_3` |               3 |                  off | Introduce only periodic rollout-policy staleness.                 |
+| `temperature_0p6`          |               1 |                  0.6 | Change only the behavior-policy temperature.                      |
+| `temperature_1p2`          |               1 |                  1.2 | Change only the behavior-policy temperature.                      |
+
+P3O and GRPO launchers for the same scenario share all non-algorithm
+configuration. Temperature scenarios preserve `top_p`, `top_k`, response
+limits, and evaluation sampling settings.
+
+## Running
+
+```bash
+bash examples/algorithms/p3o/run_p3o_on_policy_a100x4.sh
+bash examples/algorithms/p3o/run_grpo_on_policy_a100x4.sh
+bash examples/algorithms/p3o/run_p3o_periodic_sync_interval_3_a100x4.sh
+bash examples/algorithms/p3o/run_p3o_temperature_0p6_a100x4.sh
+```
+
+For a one-rollout check, select any scenario through the smoke wrapper:
+
+```bash
+bash examples/algorithms/p3o/run_p3o_smoke.sh p3o_temperature_1p2
+```
+
+Use `P3O_DRY_RUN=1` to print the resolved training arguments without checking
+assets or submitting a Ray job.
+
+## Policy-age metric
+
+`train/p3o/rollout_policy_age_rollouts` measures the difference between the
+current rollout ID and the rollout-policy snapshot ID that generated the batch.
+Its unit is rollouts, not optimizer steps. A periodic refresh affects the next
+rollout; metrics for the batch at the refresh boundary still describe the
+snapshot that generated that batch.