feat(p3o): add ESS-adaptive policy optimization - #234
Closed
DreamEnding wants to merge 45 commits into
Closed
Conversation
…cope Implements P3O (arXiv:2605.12380): replaces PPO's fixed clip range with a one-sided cap derived from the normalized ESS of token-level importance ratios, plus an adaptive trust region weighted by (1 - ESS). The paper's Algorithm 2 and the reference implementation both compute ESS per micro-batch, which makes the cap a function of the gradient-accumulation factor. Task 40 requires that neither the micro-batch count nor the DP/CP split move the cap, so ESS is instead computed over one whole optimizer step: a no-grad stats pass accumulates S1/S2/N across the window, one all-reduce over DP x CP produces the global moments, and the resulting cap is frozen into an immutable context that every micro-batch of the training pass reads. The pre-pass replays the same iterator window, so it snapshots and restores both iterator offsets and RNG state. Configurations that would break that replay (FP8 amax history, dropout, fully-async streaming) or silently change the objective (missing rollout log-probs, per-sample-mean normalization, stacked TIS) are rejected in arguments.py rather than tolerated. The reduction covers DP x CP only: TP and PP ranks hold replicas of the same tokens' log-probs, so including them would scale N and rescale the cap. Tests cover element-wise parity against frozen reference values, invariance of the cap to token partitioning, the DP/CP reduction scope, replay guards, and the config gates. The distributed matrix and end-to-end convergence runs need multi-GPU and are not exercised here.
CI runs the pre-commit action, which includes docformatter with --wrap-descriptions 79. Six P3O files failed that hook while main is fully clean, so the job would have gone red on style alone. Summary lines that the tool would have split mid-hyphen (producing "all-\nreduce", "log-\nprobs") are reworded to fit in one line instead; the remaining changes are plain paragraph rewraps applied by the tool. No code, formula, or golden value is touched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The P3O test modules imported `relax.backends.megatron.{loss,model,p3o_step}`
and `relax.core.registry`, all of which pull in `megatron.core` at module
scope. CI installs no megatron package, so collection raised
ModuleNotFoundError -- and since CI runs `pytest tests/ -x`, that aborted the
entire suite rather than skipping a few modules.
Add a shared `_megatron_stub.stubbed_megatron_modules()` context manager that
resolves `megatron.*` to synthetic MagicMock-backed modules for the duration of
the import, then restores the prior `sys.modules` state. It is a no-op when
megatron is genuinely installed, so a GPU machine still exercises the
production import path unchanged.
The tested logic is pure tensor math plus collectives and touches no megatron
symbol at call time, so the assertions now run in CI instead of silently
disappearing.
- `test_p3o_model_step.py` defers its skip rather than using
`allow_module_level=True`: `model.py` also needs `transfer_queue`, whose CI
stub has no submodules, but the AST guard test must run everywhere.
- `test_p3o_replay.py` stubs only
`megatron.core.tensor_parallel.random`, the single module
`preserved_rng_state()` imports lazily, and pre-seeds
`get_cuda_rng_tracker` so the existing monkeypatch patches rather than
invents it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`_dry_run()` spawned the launch scripts via a bare `bash` argv[0]. On Windows that resolves from `System32` before `PATH`, and `System32\bash.exe` is the WSL launcher, which runs in a separate filesystem namespace and cannot open a `D:\...` script path -- so the dry run failed for reasons unrelated to the scripts under test. Prefer an explicit Git-for-Windows bash, fall back to the usual POSIX locations, and skip rather than fail when no usable shell exists. Also decode subprocess output as UTF-8 with `errors="replace"`, since the default locale codec on Windows mangles the scripts' non-ASCII output. Behavior on Linux CI is unchanged: `/bin/bash` is found and used as before. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The optimizer-step ESS pre-pass forwarded CP-split `tokens` while
train_one_step forwards `unsplit_tokens` for VL models, so the frozen
adaptive cap was derived from a token layout the gradient pass never
evaluated. Mirror model.py's `needs_unsplit` selection, the thd bridge
packed_seq_params, and the dynamic-CP `pg_collection.cp` swap exactly.
Add `compute_p3o_sufficient_stats_unchecked`, which reports non-finite
ratios as a device-resident float64 flag instead of testing on the host.
The pre-pass now reduces that flag alongside S1/S2/N through the step's
single all-reduce, replacing one GPU-CPU sync per micro-batch. This also
removes the try/except that string-matched the error message, which
could re-raise mid-schedule on a subset of ranks and deadlock the rest
at the next collective.
Replace bare asserts in P3O argument validation and the rollout_log_probs
check with ValueError: `python -O` strips asserts, letting a misconfigured
run execute the wrong algorithm silently.
Raise instead of falling back to `device="cpu"` in get_cp_local_valid_mask
when both chunks and loss_masks are empty; every other path returns a GPU
tensor, so the fallback surfaced as an opaque device mismatch downstream.
Document the KL-proxy clamp in terms of BEHAVIOR_KL_EXP_CLAMP with its
saturation behavior, and the external-mutation hazard in
preserved_iterator_positions.
Unrun here (no ray/megatron locally):
pytest tests/backends/megatron/test_p3o_step.py \
tests/backends/megatron/test_p3o_loss.py \
tests/backends/megatron/test_p3o_distributed.py \
tests/utils/test_p3o_registry.py \
tests/components/test_p3o_advantages.py -q
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Auto-fixed by pre-commit hooks: - end-of-file-fixer: trailing newlines on docker patches and tooling stubs - docformatter: line-wrap docstrings in test_p3o_loss.py and test_configs.py No logic changes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add three TensorBoard metrics to track policy lag: - train/actor_optimizer_step - train/rollout_policy_snapshot_step - train/p3o/rollout_policy_lag_steps This enables verification of the fixed-lag rollout policy mechanism without changing the P3O algorithm logic. Addresses VERIFICATION_PLAN P1. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Add return_schedule_plan parameter to ESS pre-pass forward_step for compatibility with Megatron combined 1F1B scheduler interface - Re-export P3O public API from ppo_utils for unified namespace access - Extend test_p3o_step.py with 4 new tests covering: * Plain text forward kwargs * VL unsplit forward kwargs * VL thd bridge packed_seq_params * Dynamic CP group switching All tests pass (7 passed). Part of P3O core implementation refinement. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: DreamEnding <63937131+DreamEnding@users.noreply.github.com>
Add P3O (Adaptive Policy Optimization) with ESS-based adaptive clipping: Core algorithm implementation: - p3o_utils.py: ESS computation, sufficient statistics, and objective - p3o_step.py: two-pass optimizer-step ESS scope (stats pass + train pass) - loss.py: p3o branch with behavior KL proxy and detached adaptive cap Algorithm properties: - ESS scope: one optimizer step (partition-invariant) - Behavior policy: actual rollout sampling logprobs - Adaptive cap: min(ratio, ESS), fully detached - KL term: (1-ESS) * selected-token proxy (not full-vocabulary KL) - Replay mechanism: deterministic micro-batch iterator with frozen RNG This implementation deliberately deviates from the paper's per-micro-batch ESS to ensure partition invariance, as required by task specification.
Framework integration: - model.py: two-pass lifecycle (ESS pre-pass + train pass with frozen cap) - data.py: replay-aware iterator interface - advantages.py: register 'p3o' advantage estimator - registry.py: add P3O to algorithm registry - arguments.py: add --advantage-estimator=p3o, --p3o-ess-epsilon - ppo_utils.py: narrow re-export of compute_ppo_loss for compatibility - utils.py: add policy_entropy calculation helper The two-pass design ensures partition invariance: stats reduce over all micro-batches once before the train pass applies a uniform cap.
Context parallel compatibility: - cp_utils.py: replace assert with ValueError for production safety - cp_utils.py: ensure CP-aware logits/logprobs handling in P3O code paths Actor integration: - actor.py: rollout policy periodic sync with configurable interval - actor.py: track rollout_policy_snapshot_rollout for observability All assert statements replaced with explicit ValueError raises to meet production code safety requirements.
Add rollout_policy_lag.py module for tracking policy freshness: - compute_rollout_policy_age_rollouts: measure staleness in rollout units - Logged as train/p3o/rollout_policy_age_rollouts metric - Tracks drift between training policy and rollout policy snapshots This metric is critical for understanding P3O behavior under periodic synchronization (update_weights_interval > 1).
Unit tests: - test_p3o_utils.py: ESS computation, golden oracle validation - test_p3o_replay.py: iterator replay determinism - test_p3o_step.py: two-pass ESS lifecycle - test_p3o_loss.py: objective formula correctness - test_rollout_policy_lag.py: age computation logic Integration tests: - test_p3o_partition_invariance.py: DP/CP split invariance - test_p3o_distributed.py: multi-GPU correctness - test_p3o_on_policy.py: gradient agreement with PPO/GRPO at ESS≈1 - test_p3o_observability.py: metric logging coverage - test_p3o_model_step.py: full train_one_step lifecycle Component tests: - test_p3o_advantages.py: advantage estimator registration - test_p3o_registry.py: algorithm registry - test_p3o_arguments.py: CLI argument parsing Example tests: - test_configs.py: launcher script validation - test_rollout.py: rollout.py environment variable handling All tests are CPU-safe and CI-compatible.
A100×4 experiment suite: - common_a100x4.sh: shared configuration (model, training, Ray, observability) - rollout.py: behavior temperature parameterization via environment variables Baseline experiments: - run_p3o_on_policy_a100x4.sh: P3O with on-policy behavior (interval=1) - run_grpo_on_policy_a100x4.sh: GRPO baseline Mismatch experiments (behavior != training): Temperature mismatch (single-variable): - run_p3o_temperature_0p6_a100x4.sh: P3O with T_behavior=0.6, T_train=1.0 - run_p3o_temperature_1p2_a100x4.sh: P3O with T_behavior=1.2, T_train=1.0 - run_grpo_temperature_0p6_a100x4.sh: GRPO baseline T=0.6 - run_grpo_temperature_1p2_a100x4.sh: GRPO baseline T=1.2 Staleness mismatch (periodic sync): - run_p3o_periodic_sync_interval_3_a100x4.sh: update_weights_interval=3 - run_grpo_periodic_sync_interval_3_a100x4.sh: GRPO baseline interval=3 Smoke test: - run_p3o_smoke.sh: single-GPU 3-step quick validation All launchers explicitly set temperature and top_p values as required by task specification. Uninitialized shell variables fixed.
Apply automated formatting fixes from pre-commit hooks: - Wrap long error messages to fit 119-char line limit - Adjust docstring line breaks for readability No functional changes.
# Conflicts: # relax/backends/megatron/actor.py
# Conflicts: # relax/backends/megatron/loss.py # relax/utils/arguments.py
# 🔩 Chore ## Resolve the diverged P3O branch - Adopt the audited fork implementation for all P3O conflicts - Include the fork mainline updates already merged into the remote branch - Drop local-only Task40 analysis, probes, and fixed-lag launchers - Restore AGENTS.md to the remote version
…e-cap # Conflicts: # tests/backends/sglang/test_router_registration.py
- Add missing copyright header to cp_utils.py - Fix Python syntax error in common_a100x4.sh (import statements on separate lines) - Add type annotations to _validate_p3o_args, forward_step, and collect functions - Clarify metric semantics with inline comments: * behavior_kl_proxy: sampled-token k3 proxy (1-ESS), not full-vocabulary KL * cap_fraction: adaptive cap utilization, distinct from PPO's clip_fraction - Document intentional duplication in p3o_step.py forward input preparation to prevent silent drift between stats pass and training pass These fixes address code standards violations and reduce potential confusion around P3O-specific metrics vs standard PPO terminology. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The test extracts _validate_p3o_args via AST to avoid importing the full Megatron/Ray chain. Changed type annotation from Any to argparse.Namespace (more precise) and injected argparse into the extracted module namespace so the annotation resolves at exec time. Fixes CI collection error: NameError: name 'argparse' is not defined
Contributor
Author
This was referenced Aug 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
This PR introduces P3O (Policy-on Policy-off Policy Optimization) in Relax:
p3owith the Megatron loss path, GRPO-style advantages, argument validation, and the algorithm registry;T=0.6/1.2), periodic-sync mismatch, and GRPO comparison launchers;Why
Fixed PPO/GRPO clip ranges cannot adapt to rollout-policy mismatch caused by sampling changes or stale rollouts. P3O uses the batch policy-ratio ESS to control both the score-function weight and the behavioral-policy regularizer without adding a new clipping hyperparameter.
Relates to #233.
Official task: https://github.com/redai-infra/community/blob/main/contributor-program/2026-cohort-1/official-task.md#40-p3o--adaptive-policy-optimization高级
How
The Megatron actor performs a no-grad ESS pre-pass over the same replayable micro-batches used by the training pass. Each pipeline-last rank accumulates detached FP64 sufficient statistics (
sum(ratio),sum(ratio^2), and valid-token count), reduces them over DP x CP, and broadcasts the result over PP. The resulting optimizer-step context is frozen for every micro-batch in the train pass.The per-token objective is:
-stop_gradient(min(ratio, ESS)) * log_prob * advantage;(1 - ESS) * KL_proxy(policy || behavior).P3O reuses Relax's existing GRPO-style group-relative advantages. Temperature mismatch changes only the behavior-policy sampling temperature; periodic-sync mismatch changes only the rollout-policy refresh interval. Run metadata records the commit, branch, dirty state, resolved arguments, job status, and exit code.
Testing
Current HEAD:
942ca60197fb760ef2b63da92e46a27c5b159091.Focused tests runnable in the local CPU environment:
Result:
158 passed, 2 warnings.pre-commit run --files $(git diff --name-only upstream/main...HEAD) --show-diff-on-failure git diff --check upstream/main...HEADResult: all applicable hooks passed;
git diff --checkpassed.Known validation gaps:
tests/utils/test_p3o_registry.pyis not collectable locally becauserayis not installed;tests/components/test_p3o_advantages.pyis not collectable locally becausetransfer_queueis not installed;full
pytest tests/and Linuxpre-commit run --all-fileshave not passed on the current HEAD;real Megatron/NCCL A100x4 smoke, DP/CP/PP invariance, three-seed paper-aligned experiments, and FeynRL end-to-end comparison remain to be run.
pre-commit run --all-filespassesTests pass (
pytest tests/)New tests added (if applicable)
Documentation updated (if applicable)
Type of Change
Screenshots / Logs
No screenshots are included. The focused test and static-check results are listed above. A100x4 training curves, three-seed aggregates, throughput, and peak-memory evidence will be added after the formal validation matrix completes.