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

[DreamEnding]

Summary

• Closes 【Task 40】P3O / Adaptive Policy Optimization 论文复现进展 #233
• 解决了什么问题：新增 P3O（Policy-on Policy-off Policy Optimization），用 rollout policy-ratio 的 ESS 自适应控制 score-function 权重与行为策略正则项，处理采样参数变化和 rollout 陈旧性造成的 rollout-policy mismatch。
• 为什么采用该方案：P3O 不增加新的 clipping 超参数，复用 Relax 的 GRPO 风格 group-relative advantage 与 Megatron 训练路径；同时将统计 pass 与训练 pass 的前向输入收敛为同一 helper，避免多模态/THD 路径在后续演进中漂移。

Changes

• 文件 / 核心函数 / 行为变化：
  • relax/utils/training/p3o_utils.py、relax/backends/megatron/loss.py 实现归一化 ESS、自适应 cap、P3O policy loss 与指标；
  • relax/backends/megatron/p3o_step.py 实现 replayable no-grad ESS pre-pass；actor.py、rollout_policy_lag.py 处理 rollout-policy age 与 step context；
  • relax/backends/megatron/data.py::build_rl_forward_kwargs 被训练 pass 和 P3O 统计 pass 共用，统一普通 token、unsplit token、THD VL bridge 和 multimodal kwargs 选择；
  • relax/backends/megatron/cp_utils.py 明确动态 CP 合并过程中的长度元数据命名；
  • 新增 A100×4 的 on-policy、温度失配（T=0.6/1.2）、periodic-sync 失配及 P3O/GRPO 对照配方，并补齐中英文算法文档和中文配方 README。
• CLI、配置或兼容性变化：
  • 新增 --advantage-estimator p3o；P3O 要求 --use-rollout-logprobs 与 --calculate-per-token-loss；
  • --p3o-kl-mode 的生产 choices 收敛为 proxy 与 proxy_safe；全词表 exact 仅保留为纯验证辅助函数，不再是 CLI 模式；
  • --p3o-ess-scope 支持 micro-batch 与 step；P3O 与 --use-opd 互斥；
  • 本 PR 不含 SGLang deterministic-sampler 修复，该修复已拆分至 fix(sglang): clamp deterministic sampler endpoints #271。

Verification

• HEAD 3313cf569407ad3f75ec29174e2c0c4814c4386e。

• 可复制命令：

  python -m pytest -q \
    tests/backends/megatron/test_data_vpp.py \
    tests/backends/megatron/test_p3o_step.py \
    tests/utils/test_p3o_arguments.py \
    tests/examples/algorithms/p3o/test_configs.py \
    tests/examples/algorithms/p3o/test_rollout.py
  
  python -m pytest -q \
    tests/backends/megatron/test_p3o_cp_metadata.py \
    tests/backends/megatron/test_p3o_step.py
  
  git diff --check upstream/main...HEAD

• 单元/集成/端到端测试结果：第一组 101 passed, 2 warnings；CP 元数据回归 24 passed, 2 warnings；git diff --check 通过。PP>1 / CP>1 多节点 GPU 集成 smoke、三 seed 正式实验和 FeynRL 端到端对比未在本地运行。

• 性能 before/after（若适用）：不适用；本 PR 的目标是算法正确性和实验能力，尚未产生可比较的 A100×4 性能数据。

• 日志、曲线、profile 链接：无；A100×4 训练曲线、三 seed 聚合、吞吐与峰值显存将在正式验证矩阵完成后补充。

Risk & Rollback

• 已知限制：tests/components/test_p3o_advantages.py 需要 transfer_queue，tests/engine/rollout/test_sglang_rollout_diagnostics.py 需要 pybase64，tests/utils/test_p3o_registry.py 被当前环境的 scikit-learn/NumPy ABI 不匹配阻断；完整 pytest tests/ 未在当前环境通过。pre-commit run --all-files 会被既有 symlink/patch 文件的 end-of-file-fixer 改动中断，相关无关改动已还原。
• 风险：真实 PP × DP × CP 和 NCCL 调度尚需多节点 GPU 验证；step 作用域依赖可重放前向，相关不确定性特性已由参数校验限制。
• 关闭开关或回退方式：使用 --advantage-estimator grpo 即可不启用 P3O；紧急回退可 revert P3O 提交 8f4c368，后续可读性提交 c007b8e 与文档提交 3313cf5 可独立回退。

Checklist

• Diff 仅包含本任务必要改动
• 新增/相关测试全部通过
• 文档、脚本和默认值已更新
• 不含密钥、数据集、checkpoint 或机器隐私信息
• 已逐条回复 review comment

[Copilot]

Pull request overview

This PR adds the P3O (ESS-adaptive policy optimization) algorithm to Relax’s RL training stack, including Megatron backend loss integration, optional optimizer-step ESS pre-pass replay, and rollout-side behavior logprob alignment/masking needed to make off-policy correction robust under multimodal token sanitization and rollout staleness.

Changes:

• Introduces P3O math primitives (ESS, adaptive cap, behavior-KL proxy) plus Megatron loss dispatch + metrics for --advantage-estimator p3o.
• Adds a replayable optimizer-step ESS pre-pass (iterator position + RNG preservation) and rollout-policy snapshot/age observability for periodic-sync mismatch scenarios.
• Extends rollout payload handling with a behavior logprob pairing mask (esp. for multimodal token replacement) and adds extensive unit/distributed tests + docs/recipes.

Reviewed changes

Copilot reviewed 60 out of 60 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/utils/training/test_p3o_replay.py	Tests replay guards (iterator offsets + RNG restoration) for step-scoped ESS pre-pass.
tests/utils/test_rollout_logprob_mask.py	Tests propagation and loss-mask intersection logic for rollout behavior-logprob pairing masks.
tests/utils/test_p3o_registry.py	Validates P3O registry wiring, CLI exposure, and reward normalization parity rules.
tests/utils/test_p3o_arguments.py	Tests P3O argument validation/gates (ESS scope, KL mode, incompatibilities).
tests/utils/test_multimodal_rollout_stats.py	Verifies rollout summary records include response token IDs + behavior logprob/mask fields.
tests/utils/test_arguments_opd_teacher_colocate.py	Ensures P3O and OPD are rejected as mutually exclusive in args validation.
tests/examples/algorithms/p3o/test_rollout.py	Tests behavior-only temperature override helper for mismatch experiments.
tests/engine/rollout/test_sglang_rollout_diagnostics.py	Adds diagnostics tests for multimodal token sanitization pairing-mask behavior.
tests/components/test_p3o_advantages.py	Ensures P3O uses GRPO-style advantage path (shape/value parity).
tests/backends/megatron/test_rollout_policy_lag.py	Tests rollout-policy snapshot scheduling helpers.
tests/backends/megatron/test_p3o_step.py	Unit-tests stats synchronization + forward-kwargs parity for ESS pre-pass.
tests/backends/megatron/test_p3o_partition_invariance.py	Validates objective/grad invariance under logical partitioning.
tests/backends/megatron/test_p3o_on_policy.py	Confirms P3O degenerates correctly to on-policy policy-gradient behavior.
tests/backends/megatron/test_p3o_observability.py	Tests rollout-policy age metrics + AST checks for observability wiring.
tests/backends/megatron/test_p3o_model_step.py	Tests exception safety (dynamic CP group restoration) and guard coverage in train step.
tests/backends/megatron/test_p3o_loss.py	Tests P3O loss schema/dispatch + token-normalizer correctness.
tests/backends/megatron/test_p3o_distributed.py	Real Gloo tests for synchronized failure + stats/objective invariance.
tests/backends/megatron/test_p3o_cp_metadata.py	Fail-fast tests for CP metadata length alignment across consumers.
tests/backends/megatron/test_data_vpp.py	Adds tests for shared build_rl_forward_kwargs behavior across text/VL paths.
tests/backends/megatron/_megatron_stub.py	Adds import-time Megatron stubs to keep CPU CI tests collectible.
scripts/models/qwen3-4B.sh	Makes RoPE base configurable via env var for recipe alignment.
relax/utils/utils.py	Adds rollout logprob presence/length checks, pairing-mask handling, and P3O reward normalization rules.
relax/utils/types.py	Extends Sample with rollout_log_probs_mask to track token/logprob correspondence.
relax/utils/training/train_dump_utils.py	Adds response token IDs and behavior logprob/mask fields to rollout summary dumps.
relax/utils/training/ppo_utils.py	Defines shared estimator sets (GRPO-style advantage + group reward normalization).
relax/utils/training/p3o_utils.py	Implements P3O math: sufficient stats, step context, proxy/proxy_safe KL, token-level terms.
relax/utils/training/p3o_replay.py	Adds replay guards for RNG and iterator positions for step-scoped ESS.
relax/utils/training/data_fields.py	Adds rollout_log_probs_mask field when --use-rollout-logprobs is enabled.
relax/utils/opd/opd_utils.py	Adds explicit P3O/OPD compatibility validation and enforces it in OPD validation.
relax/utils/data/processing_utils.py	Adds multimodal token sanitization that also produces a behavior-logprob pairing mask.
relax/utils/arguments.py	Adds P3O CLI args and enforces strict validation after effective-value overrides.
relax/engine/rollout/sglang_rollout.py	Adds response token/logprob length validation, sanitization pairing mask, and mask accumulation checks.
relax/core/registry.py	Registers p3o as an algorithm mapping to GRPO-equivalent service roles.
relax/components/advantages.py	Treats P3O as GRPO-style for advantage/return computation (loss differs later).
relax/backends/megatron/rollout_policy_lag.py	Implements periodic snapshot scheduling + policy-age metrics for rollout staleness observability.
relax/backends/megatron/p3o_step.py	Implements optimizer-step ESS pre-pass, DP×CP reduce + PP broadcast, and step-context publishing.
relax/backends/megatron/model.py	Integrates step-scoped ESS pre-pass + dynamic-CP group restoration guard and logs policy-age metrics.
relax/backends/megatron/loss.py	Adds P3O loss path, step/micro-batch ESS resolution, and dispatch guard to avoid OPD hybridization.
relax/backends/megatron/data.py	Adds shared build_rl_forward_kwargs and makes DataIterator replayable via snapshot/restore.
relax/backends/megatron/cp_utils.py	Adds CP metadata length validation and introduces get_cp_local_valid_mask for exact token alignment.
relax/backends/megatron/actor.py	Switches to Relax logger, adds periodic rollout-policy snapshot mechanics, and publishes snapshot version for metrics.
examples/algorithms/README.md	Documents P3O in the supported algorithms list and describes activation requirements/constraints.
examples/algorithms/p3o/run_p3o_temperature_1p2_a100x4.sh	Adds A100×4 P3O launcher for temperature mismatch (T=1.2).
examples/algorithms/p3o/run_p3o_temperature_0p6_a100x4.sh	Adds A100×4 P3O launcher for temperature mismatch (T=0.6).
examples/algorithms/p3o/run_p3o_smoke.sh	Adds a scenario selector wrapper for one-rollout smoke runs.
examples/algorithms/p3o/run_p3o_periodic_sync_interval_3_a100x4.sh	Adds A100×4 P3O launcher for periodic-sync staleness mismatch.
examples/algorithms/p3o/run_p3o_on_policy_a100x4.sh	Adds A100×4 P3O on-policy launcher.
examples/algorithms/p3o/run_grpo_temperature_1p2_a100x4.sh	Adds paired GRPO launcher for temperature mismatch (T=1.2).
examples/algorithms/p3o/run_grpo_temperature_0p6_a100x4.sh	Adds paired GRPO launcher for temperature mismatch (T=0.6).
examples/algorithms/p3o/run_grpo_periodic_sync_interval_3_a100x4.sh	Adds paired GRPO launcher for periodic-sync mismatch.
examples/algorithms/p3o/run_grpo_on_policy_a100x4.sh	Adds paired GRPO on-policy launcher.
examples/algorithms/p3o/rollout.py	Adds behavior-only temperature override wrapper for controlled mismatch experiments.
examples/algorithms/p3o/README.md	Adds English recipe documentation and experiment contract details.
examples/algorithms/p3o/README_zh.md	Adds Chinese recipe documentation and experiment contract details.
examples/algorithms/p3o/init.py	Adds package marker/docstring for P3O recipe helpers.
docs/zh/examples/algorithms.md	Documents P3O CLI/behavior and links to recipes (Chinese).
docs/en/examples/algorithms.md	Documents P3O CLI/behavior and links to recipes (English).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[yulingma545-ctrl]

Summary

解决了什么问题：本次更新逐条回应了上一轮 review 提出的 5 个代码风险点与 S2 数值先决，并完成 P0 集群实验（Step-0 + FP32 对照、分布式边界 smoke）。

新增 7 个提交：

22635d4 fix(p3o): harden behavior policy safeguards
8679b34 test: format P3O sampling contract test
3cdf4a9 test: preserve import cache in P3O tests
0567376 test: isolate P3O Megatron stubs
7e11a49 test: wait for reward worker teardown
6c7a3d2 fix(rewards): isolate worker actor pools
cd020f1 fix(p3o): harden distributed validation

Changes

• 数值安全与分布式校验：
  • dummy micro-batch 对 micro-batch ESS 贡献全零 S1/S2/N 且 valid_mask 置零，所有 rank 仍进 collective，避免 dummy 把非零统计带入其他 real rank 的 cap（loss.py::get_p3o_context/p3o_loss_function 新增 is_dummy 路径）；
  • 全局零 valid token 时基于全局 token 数做 rank 一致 fail-fast：pipeline-last 在 DP×CP 归约、经 PP 广播到所有 stage，在 finalize_model_grads_func 前拦截，optimizer/scheduler 均不推进（model.py::_require_p3o_global_valid_tokens）；
  • S2 = sum(rho^2) 溢出/下溢加固：ratio_sq 可表示性检查（exp(500) 溢出、exp(-500) 精确为零）、local-sum 溢出检查、collective 归约后对非有限 S1/S2/N 做 all-rank 一致拒绝，杜绝静默退化为 ESS=1（p3o_utils.py、p3o_step.py）。
• Behavior-logprob 配置契约 hard-gate：
  • --rollout-top-p 必须为 1.0、--rollout-top-k 必须为 -1；拒绝 min_p/top_a/typical_p/epsilon_cutoff/eta_cutoff 截断参数；custom generation 必须声明 p3o_behavior_logprob_contract = True（arguments.py、rollout.py、sglang_rollout.py::_dispatch_generate）；
  • 拒绝 --recompute-loss-function 与 --p3o-ess-scope micro-batch 组合（checkpoint 重放会重入 ESS collective）。
• Resume / 加载语义：
  • P3O 在 update_weights_interval > 1 时从 Megatron checkpoint 精确恢复声明 unsupported（fail-fast，rollout_policy_lag.py::validate_p3o_periodic_snapshot_resume）；
  • _configure_megatron_checkpoint_loading 统一 Megatron 原生 checkpoint 直载 / HF cold-load 判定，避免直接 iter_0000000 路径被静默降级为 cold start；
  • model_provider.py：精度 flag 均关闭时显式 params_dtype = torch.float32，不继承 HF checkpoint 的降精度提示。
• ESS scope 语义：
  • --p3o-ess-scope 明确 micro-batch = 论文/参考实现兼容但拓扑相关，step = 分区不变的 adaptive cap；
  • step scope 要求 dropout=0，拒绝 FP8 与 fully-async 等非确定性 replay。
• 测试与 CI 稳定性：import-cache 保护（temporarily_stub_module）、Megatron stub 隔离（isolated_module_cache）、reward worker actor pool 隔离（匿名 actor，避免 Ray teardown 阶段重复获取同名 actor）；新增/扩充分布式、observability、sampling-contract、model-provider 等 P3O 测试。

Verification

• HEAD cd020f1。
• 单元/集成测试：工程审阅点 3313cf5 与后续加固共 101 passed + CP 元数据回归 24 passed；P3O focused 冻结点（a06f465）187 passed；P0 集群 S2 先决 48 passed、最终 focused 回归 67 passed；git diff --check 通过。
• P0 集群验证（commit 6c7a3d2，真实 Ray/Megatron/NCCL/SGLang、4×A100-40GB，实验包 task40_p0_cluster_20260814_6c7a3d2）：
  • P0-1 固定 Step-0 + FP32 对照 —— FAIL（归因重开）：同一 fixture（64 samples / 13,635 valid tokens / seed 42）下，behavior logprob max-abs 差 = 0、valid-token keys 完全一致（排除输入/replay 差异）；但 FP32 对照存在 attention-backend 混淆（BF16=flash、FP32=unfused，且经未跟踪 hook 注入精度），不是 precision-only，因此既不能否定 BF16 数值噪声假设、也不能确认生产 CP bug；分区不变性保持未闭环，冻结阈值不放宽。
  • P0-2 分布式边界 smoke —— 全部 PASS：boundary_nccl_a2/verdict.json overall_passed=true，覆盖 real+dummy、全局零 valid token、极端 ratio ±500（DP2 与 CP2）、正常未截断 rollout（64/64 samples、response 72–401、13,637 valid tokens、loss 0.00284698、normalized ESS 0.99939）。
  • 容器最终 focused 回归 67 passed, 28 warnings in 318.51s；targeted pre-commit hooks 通过。
• 未运行：当前 HEAD cd020f1 的 GPU 闭环复验（P0 覆盖到 6c7a3d2）、PP>1/CP>1 多节点 GPU smoke、三 seed 正式实验、FeynRL canonical 端到端对比（8 GPU 资源缺口）。

Risk & Rollback

• 风险：真实模型 DP/CP 分区不变性未闭环——P0-1 的 FP32 对照存在 attention-backend 混淆，须在同路径（thd+flash）FP32 或独立 per-token FP32 oracle 下重跑后定级，不放宽冻结阈值；cd020f1（P0 之后的分布式校验加固）尚无 A100 证据。
• 已知限制：完整 pytest tests/ 未在当前环境通过（test_p3o_advantages.py 需 transfer_queue、test_sglang_rollout_diagnostics.py 需 pybase64、test_p3o_registry.py 被 scikit-learn/NumPy ABI 不匹配阻断）；pre-commit run --all-files 会被既有 symlink/patch 文件的 end-of-file-fixer 改动中断，无关改动已还原。
• 关闭开关或回退方式：使用 --advantage-estimator grpo 即可不启用 P3O；紧急回退可 revert 8f4c368，c007b8e、3313cf5、22635d4、cd020f1 可独立回退。

Checklist

• Diff 仅包含本任务必要改动
• 新增/相关测试通过（CPU 单测 + P0-2 真实集群边界）
• 文档、脚本和默认值已更新
• 不含密钥、数据集、checkpoint 或机器隐私信息
• 已逐条回复 review comment（5 个代码风险点 + S2 先决均已修复并有证据）