From 8f4c368ddcf468920028be1d125b40a9802c7183 Mon Sep 17 00:00:00 2001 From: Chream Date: Thu, 13 Aug 2026 18:53:11 +0800 Subject: [PATCH 01/30] feat(p3o): add ESS-adaptive policy optimization --- docs/en/examples/algorithms.md | 36 + docs/zh/examples/algorithms.md | 33 + examples/algorithms/README.md | 26 +- examples/algorithms/p3o/README.md | 158 ++++ examples/algorithms/p3o/README_zh.md | 135 ++++ examples/algorithms/p3o/__init__.py | 3 + examples/algorithms/p3o/common_a100x4.sh | 498 ++++++++++++ examples/algorithms/p3o/rollout.py | 54 ++ .../p3o/run_grpo_on_policy_a100x4.sh | 11 + ...un_grpo_periodic_sync_interval_3_a100x4.sh | 11 + .../p3o/run_grpo_temperature_0p6_a100x4.sh | 12 + .../p3o/run_grpo_temperature_1p2_a100x4.sh | 12 + .../p3o/run_p3o_on_policy_a100x4.sh | 11 + ...run_p3o_periodic_sync_interval_3_a100x4.sh | 11 + examples/algorithms/p3o/run_p3o_smoke.sh | 62 ++ .../p3o/run_p3o_temperature_0p6_a100x4.sh | 12 + .../p3o/run_p3o_temperature_1p2_a100x4.sh | 12 + relax/backends/megatron/actor.py | 83 +- relax/backends/megatron/cp_utils.py | 149 +++- relax/backends/megatron/data.py | 58 ++ relax/backends/megatron/loss.py | 252 +++++- relax/backends/megatron/model.py | 117 +-- relax/backends/megatron/p3o_step.py | 289 +++++++ relax/backends/megatron/rollout_policy_lag.py | 102 +++ relax/components/advantages.py | 5 +- relax/core/registry.py | 7 + relax/engine/rollout/sglang_rollout.py | 64 +- relax/utils/arguments.py | 118 +++ relax/utils/data/processing_utils.py | 45 ++ relax/utils/opd/opd_utils.py | 18 + relax/utils/training/data_fields.py | 2 + relax/utils/training/p3o_replay.py | 96 +++ relax/utils/training/p3o_utils.py | 474 ++++++++++++ relax/utils/training/ppo_utils.py | 4 + relax/utils/training/train_dump_utils.py | 6 + relax/utils/types.py | 1 + relax/utils/utils.py | 70 +- scripts/models/qwen3-4B.sh | 2 +- tests/backends/megatron/_megatron_stub.py | 110 +++ tests/backends/megatron/test_data_vpp.py | 72 ++ .../backends/megatron/test_p3o_cp_metadata.py | 44 ++ .../backends/megatron/test_p3o_distributed.py | 212 ++++++ tests/backends/megatron/test_p3o_loss.py | 173 +++++ .../backends/megatron/test_p3o_model_step.py | 69 ++ .../megatron/test_p3o_observability.py | 141 ++++ tests/backends/megatron/test_p3o_on_policy.py | 70 ++ .../megatron/test_p3o_partition_invariance.py | 110 +++ tests/backends/megatron/test_p3o_step.py | 527 +++++++++++++ .../megatron/test_rollout_policy_lag.py | 48 ++ tests/components/test_p3o_advantages.py | 62 ++ .../test_sglang_rollout_diagnostics.py | 44 +- tests/examples/algorithms/p3o/test_configs.py | 719 ++++++++++++++++++ tests/examples/algorithms/p3o/test_rollout.py | 91 +++ .../test_arguments_opd_teacher_colocate.py | 8 + tests/utils/test_multimodal_rollout_stats.py | 5 + tests/utils/test_p3o_arguments.py | 209 +++++ tests/utils/test_p3o_registry.py | 120 +++ tests/utils/test_rollout_logprob_mask.py | 84 ++ tests/utils/training/test_p3o_replay.py | 194 +++++ tests/utils/training/test_p3o_utils.py | 501 ++++++++++++ 60 files changed, 6528 insertions(+), 144 deletions(-) create mode 100644 examples/algorithms/p3o/README.md create mode 100644 examples/algorithms/p3o/README_zh.md create mode 100644 examples/algorithms/p3o/__init__.py create mode 100755 examples/algorithms/p3o/common_a100x4.sh create mode 100644 examples/algorithms/p3o/rollout.py create mode 100755 examples/algorithms/p3o/run_grpo_on_policy_a100x4.sh create mode 100755 examples/algorithms/p3o/run_grpo_periodic_sync_interval_3_a100x4.sh create mode 100755 examples/algorithms/p3o/run_grpo_temperature_0p6_a100x4.sh create mode 100755 examples/algorithms/p3o/run_grpo_temperature_1p2_a100x4.sh create mode 100755 examples/algorithms/p3o/run_p3o_on_policy_a100x4.sh create mode 100755 examples/algorithms/p3o/run_p3o_periodic_sync_interval_3_a100x4.sh create mode 100755 examples/algorithms/p3o/run_p3o_smoke.sh create mode 100755 examples/algorithms/p3o/run_p3o_temperature_0p6_a100x4.sh create mode 100755 examples/algorithms/p3o/run_p3o_temperature_1p2_a100x4.sh create mode 100644 relax/backends/megatron/p3o_step.py create mode 100644 relax/backends/megatron/rollout_policy_lag.py create mode 100644 relax/utils/training/p3o_replay.py create mode 100644 relax/utils/training/p3o_utils.py create mode 100644 tests/backends/megatron/_megatron_stub.py create mode 100644 tests/backends/megatron/test_p3o_cp_metadata.py create mode 100644 tests/backends/megatron/test_p3o_distributed.py create mode 100644 tests/backends/megatron/test_p3o_loss.py create mode 100644 tests/backends/megatron/test_p3o_model_step.py create mode 100644 tests/backends/megatron/test_p3o_observability.py create mode 100644 tests/backends/megatron/test_p3o_on_policy.py create mode 100644 tests/backends/megatron/test_p3o_partition_invariance.py create mode 100644 tests/backends/megatron/test_p3o_step.py create mode 100644 tests/backends/megatron/test_rollout_policy_lag.py create mode 100644 tests/components/test_p3o_advantages.py create mode 100644 tests/examples/algorithms/p3o/test_configs.py create mode 100644 tests/examples/algorithms/p3o/test_rollout.py create mode 100644 tests/utils/test_p3o_arguments.py create mode 100644 tests/utils/test_p3o_registry.py create mode 100644 tests/utils/test_rollout_logprob_mask.py create mode 100644 tests/utils/training/test_p3o_replay.py create mode 100644 tests/utils/training/test_p3o_utils.py diff --git a/docs/en/examples/algorithms.md b/docs/en/examples/algorithms.md index a07151f45..28e0b5804 100644 --- a/docs/en/examples/algorithms.md +++ b/docs/en/examples/algorithms.md @@ -207,11 +207,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 | @@ -224,5 +259,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) diff --git a/docs/zh/examples/algorithms.md b/docs/zh/examples/algorithms.md index 2f3ed8b2c..a059d4ab3 100644 --- a/docs/zh/examples/algorithms.md +++ b/docs/zh/examples/algorithms.md @@ -204,11 +204,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 | @@ -221,5 +253,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) diff --git a/examples/algorithms/README.md b/examples/algorithms/README.md index 5f1fdea50..1f29197ed 100644 --- a/examples/algorithms/README.md +++ b/examples/algorithms/README.md @@ -8,15 +8,16 @@ Relax 框架集成了多种策略梯度算法，均通过 `--advantage-estimator ## 支持的算法 -| 算法 | 启用参数 | 推荐场景 | -| ------------------------ | ---------------------------------------------------- | --------------------------- | -| **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 | -| **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` | 默认、大多数场景 | +| **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 信任域 | ## 选择建议 @@ -34,6 +35,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 全局归一化 diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md new file mode 100644 index 000000000..b196c089c --- /dev/null +++ b/examples/algorithms/p3o/README.md @@ -0,0 +1,158 @@ +# 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 `` 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. + +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. diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md new file mode 100644 index 000000000..fc1fb53cf --- /dev/null +++ b/examples/algorithms/p3o/README_zh.md @@ -0,0 +1,135 @@ +# P3O A100×4 配置 + +English version: [README.md](README.md)。 + +这些启动脚本会在相同的 on-policy 与可控 rollout-mismatch 场景下比较 P3O 和 +GRPO。它们面向 4 张 colocate GPU，并通过 Ray Jobs 提交训练驱动。 + +## 所需环境 + +在非 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 +``` + +`formal` 模式需要 `P3O_EVAL_DATA`；`smoke` 模式可不设置。模型、训练数据和 +Megatron 路径必须在提交 Ray job 前存在。每次运行都会在 `P3O_OUTPUT_ROOT` 下 +保存解析后的参数、命令、Git 身份、日志、Ray 状态、退出码和逐步 rollout JSONL。 +只有外部证据布局需要不同的原始 rollout 目录时，才设置 +`P3O_ROLLOUT_RESULT_DIR`；最终路径会记录在 `run_identity.env`。 + +Ray job runtime 会显式禁用继承的 HTTP proxy。SGLang 通过节点本地 IP 做健康检查 +和 worker 注册；将宿主机 proxy 变量传入 Ray worker 可能使健康引擎无法通过启动 +屏障。 + +`formal` 模式默认使用 DeepScaleR 的 `problem`/`answer` 字段；`smoke` 模式默认 +使用常见的 `question`/`answer` schema。若数据集使用其他字段，显式设置 +`P3O_INPUT_KEY` 和 `P3O_LABEL_KEY`；最终取值会记录在 `run_identity.env`。 + +`formal` 模式默认使用 `deepscaler` 规则 verifier，读取 Qwen-Thinking 的 +`` 后缀和最终 `\\boxed{...}` 答案。`smoke` 模式保留面向旧 GSM8K +资产的 `mopd` 默认值。若 smoke 使用 DeepScaleR 或其他奖励契约，请显式设置 +`P3O_RM_TYPE`；最终 reward type 会记录在 `run_identity.env`。 + +正式评估默认使用 `deepscaler` 数据集名称、每个 prompt 16 个样本、4096 token +response 上限、温度 1.0 和 top-p 0.95。受限资源实验可设置 `P3O_EVAL_NAME`、 +`P3O_EVAL_N_SAMPLES`、`P3O_EVAL_MAX_RESPONSE_LEN`、`P3O_EVAL_TEMPERATURE` 和 +`P3O_EVAL_TOP_P`。这些值只影响评估，且必须在配对算法间保持一致。 + +默认 `P3O_ROLLOUT_SHUFFLE=1` 保持普通训练行为。只有在使用预先物化的固定 prompt +调度进行配对证据时才设置为 `0`；记录该设置，以免将 shuffled run 误认为固定比较。 + +当配对实验需要 P3O 和 GRPO 共享每个样本的采样 seed 时，设置 +`P3O_DETERMINISTIC_INFERENCE=1`。解析后的开关会记录在 run identity。该开关只 +控制采样随机性；首个更新之后，即使 seed 相同，不同策略权重也应产生不同响应。 + +`formal` 模式加载 `scripts/models/qwen3-4B.sh`，目标为 +Qwen3-4B-Thinking-2507；`smoke` 模式加载 `scripts/models/qwen3-0.6B.sh`。仅在 +有意验证另一兼容模型时设置 `P3O_MODEL_CONFIG`，解析后的路径会记录在 +`run_identity.env`。正式启动器会将通用 4B 脚本的 RoPE base 覆盖为 `5000000`， +以匹配该 checkpoint 的 `config.json`；smoke 保持 `1000000`。可通过 +`P3O_MODEL_ROTARY_BASE` 有意设置兼容 override，取值同样会记录。 + +## 当前 P3O 契约 + +正式 P3O 路径使用 `--p3o-ess-scope micro-batch`、 +`--p3o-kl-mode proxy_safe` 和 `--clip-low/--clip-high 0.2` 监控边距。 +`proxy_safe` 的前向值与 FeynRL 兼容的 sampled-token proxy 相同，只修正极端负 +log-ratio 的梯度。`exact` 仅能通过纯全词表验证辅助函数调用，并不是命令行模式， +因为 rollout 数据保存的是已选 token 的行为策略 log-probability，而非完整行为 +logits。 + +P3O 使用专用 policy-loss dispatch，并与 `--use-opd` 互斥。OPD teacher loss、 +OPD advantage replacement 或 OPD-only reward 都会形成未经验证的混合目标。 +reward/verifier 名称 `P3O_RM_TYPE=mopd` 与 `--use-opd` 训练功能无关，兼容数据集 +仍可使用。 + +正式默认值为 G=16、global batch 64、micro-batch 1、rollout batch 4、response +length 4096 和 30 个 optimizer step（`--num-rollout 30`）。计划配对 seed 为 42、 +123 和 2026。smoke 使用 G=4、global batch 16、response length 128 和 1 个 +optimizer step。 + +以下环境变量可在不修改场景脚本的前提下公开这些对齐设置： + +```bash +export P3O_ESS_SCOPE=micro-batch # step 仅用于 capability/replay 验证 +export P3O_KL_MODE=proxy_safe # proxy 用于 golden parity +export P3O_CLIP_LOW=0.2 +export P3O_CLIP_HIGH=0.2 +export P3O_SEED=42 +export P3O_RM_TYPE=deepscaler # smoke 与 DeepScaleR 配对时必需 +``` + +Ray worker 默认继承正常 proxy 设置。若注入的 outbound proxy 拦截 SGLang 的节点 +本地 readiness probe，可设置 `P3O_CLEAR_RUNTIME_PROXIES=1` 以在 job runtime 内 +清除 proxy 变量。它是 opt-in，且会记录在 `run_identity.env`，因为它会同时禁用 +该 job 内所有 worker 的 proxy 访问。 + +若 A100-40GB 容量无法进行 4B pilot，应先降低 pilot response length，同时保持 +micro-batch size 为 1 并记录偏差。不要把缩小的 smoke run 视为正式证据，也不要 +默默缩减三 seed 对比。需要保持 response 的资源回退时，设置 +`P3O_ACTIVATION_RECOMPUTE=1` 启用整层统一 activation recomputation，并把 +`P3O_LOG_PROBS_CHUNK_SIZE` 设为正 token 数以对 log-probability 与 entropy +reduction 分块。这两个设置对 P3O 和 GRPO 完全一致，都会记录在 run identity； +默认 `0`/`-1` 保留原始路径。 + +## 场景 + +| 场景 | 更新间隔 | 温度 override | 含义 | +| --- | ---: | ---: | --- | +| `on_policy` | 1 | 关闭 | 以正常采样配置在每个 rollout 后同步。 | +| `periodic_sync_interval_3` | 3 | 关闭 | 只引入周期性的 rollout-policy 陈旧性。 | +| `temperature_0p6` | 1 | 0.6 | 只改变行为策略温度。 | +| `temperature_1p2` | 1 | 1.2 | 只改变行为策略温度。 | + +同一场景中的 P3O 和 GRPO 启动器共享所有非算法配置。温度场景会保持 `top_p`、 +`top_k`、response 限制和评估采样设置不变。 + +## 运行 + +```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 +``` + +一轮 rollout 检查可通过 smoke wrapper 选择任意场景： + +```bash +bash examples/algorithms/p3o/run_p3o_smoke.sh p3o_temperature_1p2 +``` + +设置 `P3O_DRY_RUN=1` 可打印解析后的训练参数，而不检查资产或提交 Ray job。 + +## Policy-age 指标 + +`train/p3o/rollout_policy_age_rollouts` 衡量当前 rollout ID 与生成当前 batch 的 +rollout-policy snapshot ID 的差值。单位是 rollout，而不是 optimizer step。周期 +刷新影响下一个 rollout；刷新边界 batch 的指标仍描述生成该 batch 的 snapshot。 diff --git a/examples/algorithms/p3o/__init__.py b/examples/algorithms/p3o/__init__.py new file mode 100644 index 000000000..5af26d92a --- /dev/null +++ b/examples/algorithms/p3o/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""P3O P3O example helpers.""" diff --git a/examples/algorithms/p3o/common_a100x4.sh b/examples/algorithms/p3o/common_a100x4.sh new file mode 100755 index 000000000..21a62649f --- /dev/null +++ b/examples/algorithms/p3o/common_a100x4.sh @@ -0,0 +1,498 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +P3O_SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +P3O_REPO_ROOT="$(cd -- "${P3O_SCRIPT_DIR}/../../.." >/dev/null 2>&1 && pwd)" +P3O_MODE="${P3O_MODE:-formal}" +if [[ -z "${P3O_MODEL_ROTARY_BASE:-}" ]]; then + if [[ "${P3O_MODE}" == "formal" ]]; then + P3O_MODEL_ROTARY_BASE="5000000" + else + P3O_MODEL_ROTARY_BASE="1000000" + fi +fi +if [[ ! "${P3O_MODEL_ROTARY_BASE}" =~ ^[1-9][0-9]*$ ]]; then + echo "P3O_MODEL_ROTARY_BASE must be a positive integer" >&2 + exit 2 +fi +MODEL_ARGS_ROTARY_BASE="${P3O_MODEL_ROTARY_BASE}" +if [[ -z "${P3O_MODEL_CONFIG:-}" ]]; then + if [[ "${P3O_MODE}" == "smoke" ]]; then + P3O_MODEL_CONFIG="${P3O_REPO_ROOT}/scripts/models/qwen3-0.6B.sh" + else + P3O_MODEL_CONFIG="${P3O_REPO_ROOT}/scripts/models/qwen3-4B.sh" + fi +fi +if [[ ! -f "${P3O_MODEL_CONFIG}" ]]; then + echo "P3O_MODEL_CONFIG does not exist: ${P3O_MODEL_CONFIG}" >&2 + exit 2 +fi +source "${P3O_MODEL_CONFIG}" + +P3O_ALGORITHM="${P3O_ALGORITHM:?set P3O_ALGORITHM to p3o or grpo}" +P3O_ENABLE_TEMPERATURE_OVERRIDE="${P3O_ENABLE_TEMPERATURE_OVERRIDE:-0}" +P3O_BEHAVIOR_TEMPERATURE="${P3O_BEHAVIOR_TEMPERATURE:-}" +P3O_MAX_STALENESS="${P3O_MAX_STALENESS:-0}" +P3O_UPDATE_WEIGHTS_INTERVAL="${P3O_UPDATE_WEIGHTS_INTERVAL:-1}" +P3O_PIPELINE_MODEL_PARALLEL_SIZE="${P3O_PIPELINE_MODEL_PARALLEL_SIZE:-1}" +P3O_SEED="${P3O_SEED:-42}" +P3O_ESS_SCOPE="${P3O_ESS_SCOPE:-micro-batch}" +P3O_KL_MODE="${P3O_KL_MODE:-proxy_safe}" +P3O_CLIP_LOW="${P3O_CLIP_LOW:-0.2}" +P3O_CLIP_HIGH="${P3O_CLIP_HIGH:-0.2}" +P3O_ACTIVATION_RECOMPUTE="${P3O_ACTIVATION_RECOMPUTE:-0}" +P3O_LOG_PROBS_CHUNK_SIZE="${P3O_LOG_PROBS_CHUNK_SIZE:--1}" +P3O_ROLLOUT_SHUFFLE="${P3O_ROLLOUT_SHUFFLE:-1}" +P3O_DETERMINISTIC_INFERENCE="${P3O_DETERMINISTIC_INFERENCE:-0}" +P3O_CLEAR_RUNTIME_PROXIES="${P3O_CLEAR_RUNTIME_PROXIES:-0}" +P3O_DRY_RUN="${P3O_DRY_RUN:-0}" +P3O_NCCL_DEBUG="${P3O_NCCL_DEBUG:-WARN}" +P3O_TORCH_DISTRIBUTED_DEBUG="${P3O_TORCH_DISTRIBUTED_DEBUG:-OFF}" +if [[ -z "${P3O_INPUT_KEY:-}" ]]; then + if [[ "${P3O_MODE}" == "formal" ]]; then + P3O_INPUT_KEY="problem" + else + P3O_INPUT_KEY="question" + fi +fi +P3O_LABEL_KEY="${P3O_LABEL_KEY:-answer}" +if [[ -z "${P3O_RM_TYPE:-}" ]]; then + if [[ "${P3O_MODE}" == "formal" ]]; then + P3O_RM_TYPE="deepscaler" + else + P3O_RM_TYPE="mopd" + fi +fi +P3O_EVAL_NAME="${P3O_EVAL_NAME:-deepscaler}" +P3O_EVAL_N_SAMPLES="${P3O_EVAL_N_SAMPLES:-16}" +P3O_EVAL_MAX_RESPONSE_LEN="${P3O_EVAL_MAX_RESPONSE_LEN:-4096}" +P3O_EVAL_TEMPERATURE="${P3O_EVAL_TEMPERATURE:-1.0}" +P3O_EVAL_TOP_P="${P3O_EVAL_TOP_P:-0.95}" + +if [[ "${P3O_DRY_RUN}" == "1" ]]; then + P3O_MODEL_DIR="${P3O_MODEL_DIR:-/dummy/model}" + P3O_TRAIN_DATA="${P3O_TRAIN_DATA:-/dummy/train.jsonl}" + P3O_EVAL_DATA="${P3O_EVAL_DATA:-/dummy/eval.jsonl}" + P3O_OUTPUT_ROOT="${P3O_OUTPUT_ROOT:-/dummy/output}" + P3O_MEGATRON_DIR="${P3O_MEGATRON_DIR:-/dummy/megatron}" +else + : "${P3O_MODEL_DIR:?P3O_MODEL_DIR must be set}" + : "${P3O_TRAIN_DATA:?P3O_TRAIN_DATA must be set}" + : "${P3O_OUTPUT_ROOT:?P3O_OUTPUT_ROOT must be set}" + : "${P3O_MEGATRON_DIR:?P3O_MEGATRON_DIR must be set}" + if [[ "${P3O_MODE}" == "formal" ]]; then + : "${P3O_EVAL_DATA:?P3O_EVAL_DATA must be set in formal mode}" + else + P3O_EVAL_DATA="${P3O_EVAL_DATA:-}" + fi +fi + +if [[ "${P3O_ROLLOUT_SHUFFLE}" != "0" && "${P3O_ROLLOUT_SHUFFLE}" != "1" ]]; then + echo "P3O_ROLLOUT_SHUFFLE must be 0 or 1" >&2 + exit 2 +fi +if [[ "${P3O_DETERMINISTIC_INFERENCE}" != "0" && "${P3O_DETERMINISTIC_INFERENCE}" != "1" ]]; then + echo "P3O_DETERMINISTIC_INFERENCE must be 0 or 1" >&2 + exit 2 +fi +if [[ "${P3O_CLEAR_RUNTIME_PROXIES}" != "0" && "${P3O_CLEAR_RUNTIME_PROXIES}" != "1" ]]; then + echo "P3O_CLEAR_RUNTIME_PROXIES must be 0 or 1" >&2 + exit 2 +fi +if [[ "${P3O_ACTIVATION_RECOMPUTE}" != "0" && "${P3O_ACTIVATION_RECOMPUTE}" != "1" ]]; then + echo "P3O_ACTIVATION_RECOMPUTE must be 0 or 1" >&2 + exit 2 +fi +if [[ "${P3O_LOG_PROBS_CHUNK_SIZE}" != "-1" && ! "${P3O_LOG_PROBS_CHUNK_SIZE}" =~ ^[1-9][0-9]*$ ]]; then + echo "P3O_LOG_PROBS_CHUNK_SIZE must be -1 or a positive integer" >&2 + exit 2 +fi +if [[ ! "${P3O_EVAL_N_SAMPLES}" =~ ^[1-9][0-9]*$ ]]; then + echo "P3O_EVAL_N_SAMPLES must be a positive integer" >&2 + exit 2 +fi +if [[ ! "${P3O_EVAL_MAX_RESPONSE_LEN}" =~ ^[1-9][0-9]*$ ]]; then + echo "P3O_EVAL_MAX_RESPONSE_LEN must be a positive integer" >&2 + exit 2 +fi + +: "${P3O_RAY_DASHBOARD:?P3O_RAY_DASHBOARD must be set}" + +if [[ "${P3O_ALGORITHM}" != "p3o" && "${P3O_ALGORITHM}" != "grpo" ]]; then + echo "Unsupported P3O_ALGORITHM=${P3O_ALGORITHM}" >&2 + exit 2 +fi +if [[ "${P3O_ENABLE_TEMPERATURE_OVERRIDE}" != "0" && "${P3O_ENABLE_TEMPERATURE_OVERRIDE}" != "1" ]]; then + echo "P3O_ENABLE_TEMPERATURE_OVERRIDE must be 0 or 1" >&2 + exit 2 +fi +if [[ "${P3O_ENABLE_TEMPERATURE_OVERRIDE}" == "1" ]]; then + python3 - "${P3O_BEHAVIOR_TEMPERATURE}" <<'PY' +import math +import sys + +try: + value = float(sys.argv[1]) +except ValueError as exc: + raise SystemExit("P3O_BEHAVIOR_TEMPERATURE must be numeric") from exc + +if not math.isfinite(value) or value <= 0.0: + raise SystemExit("P3O_BEHAVIOR_TEMPERATURE must be finite and greater than zero") +PY +fi +if [[ "${P3O_MODE}" != "formal" && "${P3O_MODE}" != "smoke" ]]; then + echo "P3O_MODE must be formal or smoke" >&2 + exit 2 +fi +if [[ "${P3O_ESS_SCOPE}" != "micro-batch" && "${P3O_ESS_SCOPE}" != "step" ]]; then + echo "P3O_ESS_SCOPE must be micro-batch or step" >&2 + exit 2 +fi +if [[ "${P3O_KL_MODE}" != "proxy" && "${P3O_KL_MODE}" != "proxy_safe" ]]; then + echo "P3O_KL_MODE must be proxy or proxy_safe for production training" >&2 + exit 2 +fi +if [[ ! "${P3O_UPDATE_WEIGHTS_INTERVAL}" =~ ^[1-9][0-9]*$ ]]; then + echo "P3O_UPDATE_WEIGHTS_INTERVAL must be a positive integer" >&2 + exit 2 +fi +if [[ "${P3O_UPDATE_WEIGHTS_INTERVAL}" != "1" && "${P3O_ENABLE_TEMPERATURE_OVERRIDE}" == "1" ]]; then + echo "periodic policy synchronization and temperature override must be tested in separate runs" >&2 + exit 2 +fi +if [[ ! "${P3O_PIPELINE_MODEL_PARALLEL_SIZE}" =~ ^[1-9][0-9]*$ ]]; then + echo "P3O_PIPELINE_MODEL_PARALLEL_SIZE must be a positive integer" >&2 + exit 2 +fi + +if [[ "${P3O_MODE}" == "formal" ]]; then + P3O_NUM_ROLLOUT="${P3O_NUM_ROLLOUT:-30}" + P3O_ROLLOUT_BATCH_SIZE="${P3O_ROLLOUT_BATCH_SIZE:-4}" + P3O_N_SAMPLES="${P3O_N_SAMPLES:-16}" + P3O_GLOBAL_BATCH_SIZE="${P3O_GLOBAL_BATCH_SIZE:-64}" + # Full-length responses make the FP32 logits conversion exceed A100-40GB at micro-batch 4. + P3O_MICRO_BATCH_SIZE="${P3O_MICRO_BATCH_SIZE:-1}" + P3O_MAX_RESPONSE_LEN="${P3O_MAX_RESPONSE_LEN:-4096}" +else + P3O_NUM_ROLLOUT="${P3O_NUM_ROLLOUT:-1}" + P3O_ROLLOUT_BATCH_SIZE="${P3O_ROLLOUT_BATCH_SIZE:-4}" + P3O_N_SAMPLES="${P3O_N_SAMPLES:-4}" + P3O_GLOBAL_BATCH_SIZE="${P3O_GLOBAL_BATCH_SIZE:-16}" + P3O_MICRO_BATCH_SIZE="${P3O_MICRO_BATCH_SIZE:-1}" + P3O_MAX_RESPONSE_LEN="${P3O_MAX_RESPONSE_LEN:-128}" +fi + +P3O_CONFIG_NAME="${P3O_ALGORITHM}_$( + if [[ "${P3O_ENABLE_TEMPERATURE_OVERRIDE}" == "1" ]]; then + echo "temperature_${P3O_BEHAVIOR_TEMPERATURE//./p}" + elif [[ "${P3O_UPDATE_WEIGHTS_INTERVAL}" != "1" ]]; then + echo "periodic_sync_interval_${P3O_UPDATE_WEIGHTS_INTERVAL}" + else + echo "on_policy" + fi +)" +if [[ "${P3O_PIPELINE_MODEL_PARALLEL_SIZE}" != "1" ]]; then + P3O_CONFIG_NAME="${P3O_CONFIG_NAME}_pp${P3O_PIPELINE_MODEL_PARALLEL_SIZE}" +fi + +P3O_build_args() { + P3O_CKPT_ARGS=( + --hf-checkpoint "${P3O_MODEL_DIR}" + --megatron-to-hf-mode bridge + --warm-hf-checkpoint-page-cache + ) + + P3O_ROLLOUT_ARGS=( + --prompt-data "${P3O_TRAIN_DATA}" + --input-key "${P3O_INPUT_KEY}" + --label-key "${P3O_LABEL_KEY}" + --apply-chat-template + --rm-type "${P3O_RM_TYPE}" + --num-rollout "${P3O_NUM_ROLLOUT}" + --rollout-batch-size "${P3O_ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${P3O_N_SAMPLES}" + --rollout-max-prompt-len 512 + --rollout-max-response-len "${P3O_MAX_RESPONSE_LEN}" + --rollout-temperature 1.0 + --rollout-top-p 1.0 + --rollout-top-k -1 + --global-batch-size "${P3O_GLOBAL_BATCH_SIZE}" + --use-rollout-logprobs + --balance-data + --log-passrate + ) + if [[ "${P3O_ROLLOUT_SHUFFLE}" == "1" ]]; then + P3O_ROLLOUT_ARGS+=(--rollout-shuffle) + fi + + P3O_PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size "${P3O_PIPELINE_MODEL_PARALLEL_SIZE}" + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --micro-batch-size "${P3O_MICRO_BATCH_SIZE}" + --calculate-per-token-loss + ) + if [[ "${P3O_ACTIVATION_RECOMPUTE}" == "1" ]]; then + P3O_PERF_ARGS+=( + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + ) + fi + if [[ "${P3O_LOG_PROBS_CHUNK_SIZE}" != "-1" ]]; then + P3O_PERF_ARGS+=(--log-probs-chunk-size "${P3O_LOG_PROBS_CHUNK_SIZE}") + fi + + P3O_OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-5 + --min-lr 0 + --lr-decay-style cosine + --lr-warmup-fraction 0.1 + --weight-decay 0.01 + --adam-beta1 0.9 + --adam-beta2 0.95 + --clip-grad 1.0 + ) + + P3O_ALGO_ARGS=( + --advantage-estimator "${P3O_ALGORITHM}" + --kl-coef 0.0 + --entropy-coef 0.0 + ) + if [[ "${P3O_ALGORITHM}" == "p3o" ]]; then + P3O_ALGO_ARGS+=( + --p3o-ess-scope "${P3O_ESS_SCOPE}" + --p3o-kl-mode "${P3O_KL_MODE}" + --clip-low "${P3O_CLIP_LOW}" + --clip-high "${P3O_CLIP_HIGH}" + ) + fi + if [[ "${P3O_ALGORITHM}" == "grpo" ]]; then + P3O_ALGO_ARGS+=(--eps-clip 0.4 --eps-clip-high 0.4) + fi + if [[ "${P3O_ENABLE_TEMPERATURE_OVERRIDE}" == "1" ]]; then + P3O_ALGO_ARGS+=(--custom-generate-function-path examples.algorithms.p3o.rollout.generate) + fi + + P3O_SGLANG_ARGS=( + --rollout-num-gpus 4 + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.70 + ) + if [[ "${P3O_DETERMINISTIC_INFERENCE}" == "1" ]]; then + P3O_SGLANG_ARGS+=(--sglang-enable-deterministic-inference) + fi + + P3O_MISC_ARGS=( + --seed "${P3O_SEED}" + --rollout-seed "${P3O_SEED}" + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --use-health-check + --use-tensorboard + --tb-project-name P3O-p3o-a100x4 + --tb-experiment-name "${P3O_CONFIG_NAME}-seed-${P3O_SEED}" + ) + + P3O_EVAL_ARGS=(--skip-eval-before-train) + if [[ "${P3O_MODE}" == "formal" ]]; then + P3O_EVAL_ARGS+=( + --eval-interval "${P3O_NUM_ROLLOUT}" + --eval-prompt-data "${P3O_EVAL_NAME}" "${P3O_EVAL_DATA}" + --n-samples-per-eval-prompt "${P3O_EVAL_N_SAMPLES}" + --eval-max-response-len "${P3O_EVAL_MAX_RESPONSE_LEN}" + --eval-temperature "${P3O_EVAL_TEMPERATURE}" + --eval-top-p "${P3O_EVAL_TOP_P}" + ) + fi + + P3O_TRAIN_ARGS=( + --resource '{"actor":[1,4],"rollout":[1,4]}' + --max-staleness "${P3O_MAX_STALENESS}" + --update-weights-interval "${P3O_UPDATE_WEIGHTS_INTERVAL}" + --num-iters-per-train-update 1 + --num-data-storage-units 1 + --colocate + "${MODEL_ARGS[@]}" + "${P3O_CKPT_ARGS[@]}" + "${P3O_ROLLOUT_ARGS[@]}" + "${P3O_PERF_ARGS[@]}" + "${P3O_OPTIMIZER_ARGS[@]}" + "${P3O_ALGO_ARGS[@]}" + "${P3O_SGLANG_ARGS[@]}" + "${P3O_EVAL_ARGS[@]}" + "${P3O_MISC_ARGS[@]}" + ) +} + +P3O_run() { + P3O_build_args + if [[ "${P3O_DRY_RUN:-0}" == "1" ]]; then + P3O_EFFECTIVE_ROLLOUT_RESULT_DIR="${P3O_ROLLOUT_RESULT_DIR:-${P3O_OUTPUT_ROOT}/rollout_results}" + P3O_TRAIN_ARGS+=(--rollout-result-dir "${P3O_EFFECTIVE_ROLLOUT_RESULT_DIR}") + printf '%s\n' "${P3O_TRAIN_ARGS[@]}" + return 0 + fi + + for required_path in "${P3O_MODEL_DIR}" "${P3O_TRAIN_DATA}" "${P3O_MEGATRON_DIR}"; do + if [[ ! -e "${required_path}" ]]; then + echo "Required P3O asset is missing: ${required_path}" >&2 + exit 2 + fi + done + if [[ "${P3O_MODE}" == "formal" && ! -e "${P3O_EVAL_DATA}" ]]; then + echo "Required P3O asset is missing: ${P3O_EVAL_DATA}" >&2 + exit 2 + fi + + P3O_RUN_ID="${P3O_RUN_ID:-$(date -u +%Y%m%dT%H%M%SZ)-$$}" + P3O_RUN_DIR="${P3O_OUTPUT_ROOT}/${P3O_CONFIG_NAME}/seed_${P3O_SEED}/${P3O_RUN_ID}" + mkdir -p "$(dirname -- "${P3O_RUN_DIR}")" + if ! mkdir "${P3O_RUN_DIR}"; then + echo "Refusing to overwrite P3O run directory: ${P3O_RUN_DIR}" >&2 + exit 2 + fi + mkdir "${P3O_RUN_DIR}/tensorboard" + P3O_EFFECTIVE_ROLLOUT_RESULT_DIR="${P3O_ROLLOUT_RESULT_DIR:-${P3O_RUN_DIR}/rollout_results}" + P3O_TRAIN_ARGS+=(--rollout-result-dir "${P3O_EFFECTIVE_ROLLOUT_RESULT_DIR}") + P3O_JOB_ID="${P3O_CONFIG_NAME}-seed-${P3O_SEED}-${P3O_RUN_ID}" + P3O_GIT_COMMIT="$(git -C "${P3O_REPO_ROOT}" rev-parse HEAD)" + P3O_GIT_BRANCH="$(git -C "${P3O_REPO_ROOT}" symbolic-ref --short -q HEAD || true)" + P3O_GIT_DIRTY=0 + if [[ -n "$(git -C "${P3O_REPO_ROOT}" status --short)" ]]; then + P3O_GIT_DIRTY=1 + fi + + printf '%s\n' "${P3O_TRAIN_ARGS[@]}" >"${P3O_RUN_DIR}/resolved_args.txt" + { + echo "GIT_COMMIT=${P3O_GIT_COMMIT}" + echo "GIT_BRANCH=${P3O_GIT_BRANCH:-DETACHED}" + echo "GIT_DIRTY=${P3O_GIT_DIRTY}" + echo "config=${P3O_CONFIG_NAME}" + echo "mode=${P3O_MODE}" + echo "seed=${P3O_SEED}" + echo "model_config=${P3O_MODEL_CONFIG}" + echo "model_rotary_base=${P3O_MODEL_ROTARY_BASE}" + echo "p3o_ess_scope=${P3O_ESS_SCOPE}" + echo "p3o_kl_mode=${P3O_KL_MODE}" + echo "clip_low=${P3O_CLIP_LOW}" + echo "clip_high=${P3O_CLIP_HIGH}" + echo "activation_recompute=${P3O_ACTIVATION_RECOMPUTE}" + echo "log_probs_chunk_size=${P3O_LOG_PROBS_CHUNK_SIZE}" + echo "max_staleness=${P3O_MAX_STALENESS}" + echo "update_weights_interval=${P3O_UPDATE_WEIGHTS_INTERVAL}" + echo "pipeline_model_parallel_size=${P3O_PIPELINE_MODEL_PARALLEL_SIZE}" + echo "nccl_debug=${P3O_NCCL_DEBUG}" + echo "torch_distributed_debug=${P3O_TORCH_DISTRIBUTED_DEBUG}" + echo "behavior_temperature=${P3O_BEHAVIOR_TEMPERATURE}" + echo "ray_job_id=${P3O_JOB_ID}" + echo "repo=${P3O_REPO_ROOT}" + echo "model=${P3O_MODEL_DIR}" + echo "train_data=${P3O_TRAIN_DATA}" + echo "input_key=${P3O_INPUT_KEY}" + echo "label_key=${P3O_LABEL_KEY}" + echo "rm_type=${P3O_RM_TYPE}" + echo "rollout_shuffle=${P3O_ROLLOUT_SHUFFLE}" + echo "deterministic_inference=${P3O_DETERMINISTIC_INFERENCE}" + echo "clear_runtime_proxies=${P3O_CLEAR_RUNTIME_PROXIES}" + echo "rollout_result_dir=${P3O_EFFECTIVE_ROLLOUT_RESULT_DIR}" + echo "eval_data=${P3O_EVAL_DATA}" + echo "eval_name=${P3O_EVAL_NAME}" + echo "eval_n_samples=${P3O_EVAL_N_SAMPLES}" + echo "eval_max_response_len=${P3O_EVAL_MAX_RESPONSE_LEN}" + echo "eval_temperature=${P3O_EVAL_TEMPERATURE}" + echo "eval_top_p=${P3O_EVAL_TOP_P}" + echo "ray_dashboard=${P3O_RAY_DASHBOARD}" + echo "started_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)" + } >"${P3O_RUN_DIR}/run_identity.env" + + P3O_RUNTIME_ENV_JSON="$( + P3O_RUNTIME_PYTHONPATH="${P3O_REPO_ROOT}:${P3O_MEGATRON_DIR}" \ + P3O_TENSORBOARD_DIR="${P3O_RUN_DIR}/tensorboard" \ + P3O_RUNTIME_ENABLE_TEMPERATURE_OVERRIDE="${P3O_ENABLE_TEMPERATURE_OVERRIDE}" \ + P3O_RUNTIME_BEHAVIOR_TEMPERATURE="${P3O_BEHAVIOR_TEMPERATURE}" \ + P3O_RUNTIME_CLEAR_PROXIES="${P3O_CLEAR_RUNTIME_PROXIES}" \ + P3O_RUNTIME_NCCL_DEBUG="${P3O_NCCL_DEBUG}" \ + P3O_RUNTIME_TORCH_DISTRIBUTED_DEBUG="${P3O_TORCH_DISTRIBUTED_DEBUG}" \ + P3O_RUNTIME_CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" \ + P3O_RUNTIME_OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" \ + P3O_RUNTIME_MKL_NUM_THREADS="${MKL_NUM_THREADS:-8}" \ + P3O_RUNTIME_OPENBLAS_NUM_THREADS="${OPENBLAS_NUM_THREADS:-8}" \ + P3O_RUNTIME_NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-0}" \ + P3O_RUNTIME_NVSHMEM_DISABLE_NCCL="${NVSHMEM_DISABLE_NCCL:-1}" \ + python3 - <<'PY' +import json +import os + +env_vars = { + "PYTHONUNBUFFERED": "1", + "PYTHONPATH": os.environ["P3O_RUNTIME_PYTHONPATH"], + "TENSORBOARD_DIR": os.environ["P3O_TENSORBOARD_DIR"], + "NCCL_DEBUG": os.environ["P3O_RUNTIME_NCCL_DEBUG"], + "TORCH_DISTRIBUTED_DEBUG": os.environ["P3O_RUNTIME_TORCH_DISTRIBUTED_DEBUG"], + "RAY_OVERRIDE_JOB_RUNTIME_ENV": "1", + "CUDA_DEVICE_MAX_CONNECTIONS": os.environ["P3O_RUNTIME_CUDA_DEVICE_MAX_CONNECTIONS"], + "OMP_NUM_THREADS": os.environ["P3O_RUNTIME_OMP_NUM_THREADS"], + "MKL_NUM_THREADS": os.environ["P3O_RUNTIME_MKL_NUM_THREADS"], + "OPENBLAS_NUM_THREADS": os.environ["P3O_RUNTIME_OPENBLAS_NUM_THREADS"], + "NCCL_NVLS_ENABLE": os.environ["P3O_RUNTIME_NCCL_NVLS_ENABLE"], + "NVSHMEM_DISABLE_NCCL": os.environ["P3O_RUNTIME_NVSHMEM_DISABLE_NCCL"], +} + +if os.environ["P3O_RUNTIME_CLEAR_PROXIES"] == "1": + # Some clusters inject an outbound proxy into the raylet. SGLang's local + # node-IP readiness probes must bypass it, but clearing worker networking is + # intentionally opt-in because other deployments require those proxies. + env_vars.update( + { + "HTTP_PROXY": "", + "HTTPS_PROXY": "", + "ALL_PROXY": "", + "http_proxy": "", + "https_proxy": "", + "all_proxy": "", + "NO_PROXY": "*", + "no_proxy": "*", + } + ) + +if os.environ["P3O_RUNTIME_ENABLE_TEMPERATURE_OVERRIDE"] == "1": + env_vars["P3O_BEHAVIOR_TEMPERATURE"] = os.environ["P3O_RUNTIME_BEHAVIOR_TEMPERATURE"] + +print(json.dumps({"env_vars": env_vars})) +PY + )" + + P3O_COMMAND=( + ray job submit + --address "${P3O_RAY_DASHBOARD}" + --submission-id "${P3O_JOB_ID}" + --runtime-env-json "${P3O_RUNTIME_ENV_JSON}" + -- + python3 -m relax.entrypoints.train + "${P3O_TRAIN_ARGS[@]}" + ) + printf '%q ' "${P3O_COMMAND[@]}" >"${P3O_RUN_DIR}/command.sh" + printf '\n' >>"${P3O_RUN_DIR}/command.sh" + + set -o pipefail + set +e + "${P3O_COMMAND[@]}" 2>&1 | tee "${P3O_RUN_DIR}/stdout_stderr.log" + P3O_EXIT_CODE=${PIPESTATUS[0]} + ray job status "${P3O_JOB_ID}" --address "${P3O_RAY_DASHBOARD}" >"${P3O_RUN_DIR}/job_status.txt" 2>&1 + P3O_STATUS_QUERY_EXIT_CODE=$? + set -e + echo "${P3O_EXIT_CODE}" >"${P3O_RUN_DIR}/exit_code.txt" + echo "${P3O_STATUS_QUERY_EXIT_CODE}" >"${P3O_RUN_DIR}/job_status_query_exit_code.txt" + echo "ended_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)" >>"${P3O_RUN_DIR}/run_identity.env" + return "${P3O_EXIT_CODE}" +} diff --git a/examples/algorithms/p3o/rollout.py b/examples/algorithms/p3o/rollout.py new file mode 100644 index 000000000..a51e7fb61 --- /dev/null +++ b/examples/algorithms/p3o/rollout.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Controlled behavior-policy sampling for the P3O mismatch experiment.""" + +import math +import os +from argparse import Namespace +from typing import Any + +from relax.utils.types import Sample + + +async def _sglang_generate(*args: Any, **kwargs: Any) -> Sample: + """Import the heavyweight rollout backend only when generation starts.""" + from relax.engine.rollout.sglang_rollout import generate + + return await generate(*args, **kwargs) + + +def _behavior_temperature() -> float: + raw_value = os.environ.get("P3O_BEHAVIOR_TEMPERATURE") + if raw_value is None: + raise ValueError("P3O_BEHAVIOR_TEMPERATURE must be set when temperature override is enabled") + try: + value = float(raw_value) + except ValueError as exc: + raise ValueError("P3O_BEHAVIOR_TEMPERATURE must be numeric") from exc + if not math.isfinite(value) or value <= 0.0: + raise ValueError("P3O_BEHAVIOR_TEMPERATURE must be finite and greater than zero") + return value + + +def behavior_sampling_params(sampling_params: dict[str, Any], *, evaluation: bool) -> dict[str, Any]: + """Return isolated sampling parameters for P3O rollout generation.""" + updated = sampling_params.copy() + if not evaluation: + updated["temperature"] = _behavior_temperature() + return updated + + +async def generate( + args: Namespace, + sample: Sample, + sampling_params: dict[str, Any], + evaluation: bool = False, +) -> Sample: + """Generate with behavior-only mismatch while preserving evaluation + settings.""" + return await _sglang_generate( + args, + sample, + behavior_sampling_params(sampling_params, evaluation=evaluation), + evaluation=evaluation, + ) diff --git a/examples/algorithms/p3o/run_grpo_on_policy_a100x4.sh b/examples/algorithms/p3o/run_grpo_on_policy_a100x4.sh new file mode 100755 index 000000000..b51084e49 --- /dev/null +++ b/examples/algorithms/p3o/run_grpo_on_policy_a100x4.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=grpo +export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 +export P3O_UPDATE_WEIGHTS_INTERVAL=1 +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_grpo_periodic_sync_interval_3_a100x4.sh b/examples/algorithms/p3o/run_grpo_periodic_sync_interval_3_a100x4.sh new file mode 100755 index 000000000..afc30105e --- /dev/null +++ b/examples/algorithms/p3o/run_grpo_periodic_sync_interval_3_a100x4.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=grpo +export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 +export P3O_UPDATE_WEIGHTS_INTERVAL="${P3O_UPDATE_WEIGHTS_INTERVAL:-3}" +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_grpo_temperature_0p6_a100x4.sh b/examples/algorithms/p3o/run_grpo_temperature_0p6_a100x4.sh new file mode 100755 index 000000000..5e2d76f2a --- /dev/null +++ b/examples/algorithms/p3o/run_grpo_temperature_0p6_a100x4.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=grpo +export P3O_UPDATE_WEIGHTS_INTERVAL=1 +export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 +export P3O_BEHAVIOR_TEMPERATURE=0.6 +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_grpo_temperature_1p2_a100x4.sh b/examples/algorithms/p3o/run_grpo_temperature_1p2_a100x4.sh new file mode 100755 index 000000000..cfa471600 --- /dev/null +++ b/examples/algorithms/p3o/run_grpo_temperature_1p2_a100x4.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=grpo +export P3O_UPDATE_WEIGHTS_INTERVAL=1 +export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 +export P3O_BEHAVIOR_TEMPERATURE=1.2 +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_p3o_on_policy_a100x4.sh b/examples/algorithms/p3o/run_p3o_on_policy_a100x4.sh new file mode 100755 index 000000000..0083e6be5 --- /dev/null +++ b/examples/algorithms/p3o/run_p3o_on_policy_a100x4.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=p3o +export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 +export P3O_UPDATE_WEIGHTS_INTERVAL=1 +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_p3o_periodic_sync_interval_3_a100x4.sh b/examples/algorithms/p3o/run_p3o_periodic_sync_interval_3_a100x4.sh new file mode 100755 index 000000000..aed6c1eca --- /dev/null +++ b/examples/algorithms/p3o/run_p3o_periodic_sync_interval_3_a100x4.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=p3o +export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 +export P3O_UPDATE_WEIGHTS_INTERVAL="${P3O_UPDATE_WEIGHTS_INTERVAL:-3}" +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_p3o_smoke.sh b/examples/algorithms/p3o/run_p3o_smoke.sh new file mode 100755 index 000000000..164d996ba --- /dev/null +++ b/examples/algorithms/p3o/run_p3o_smoke.sh @@ -0,0 +1,62 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +CONFIG="${1:-p3o_on_policy}" +case "${CONFIG}" in + p3o_on_policy) + export P3O_ALGORITHM=p3o + export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + ;; + grpo_on_policy) + export P3O_ALGORITHM=grpo + export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + ;; + p3o_temperature_0p6) + export P3O_ALGORITHM=p3o + export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 + export P3O_BEHAVIOR_TEMPERATURE=0.6 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + ;; + grpo_temperature_0p6) + export P3O_ALGORITHM=grpo + export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 + export P3O_BEHAVIOR_TEMPERATURE=0.6 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + ;; + p3o_temperature_1p2) + export P3O_ALGORITHM=p3o + export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 + export P3O_BEHAVIOR_TEMPERATURE=1.2 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + ;; + grpo_temperature_1p2) + export P3O_ALGORITHM=grpo + export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 + export P3O_BEHAVIOR_TEMPERATURE=1.2 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + ;; + p3o_periodic_sync_interval_3) + export P3O_ALGORITHM=p3o + export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 + export P3O_UPDATE_WEIGHTS_INTERVAL=3 + ;; + grpo_periodic_sync_interval_3) + export P3O_ALGORITHM=grpo + export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 + export P3O_UPDATE_WEIGHTS_INTERVAL=3 + ;; + *) + echo "Unknown smoke config: ${CONFIG}" >&2 + exit 2 + ;; +esac + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_MODE=smoke +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_p3o_temperature_0p6_a100x4.sh b/examples/algorithms/p3o/run_p3o_temperature_0p6_a100x4.sh new file mode 100755 index 000000000..779b1b725 --- /dev/null +++ b/examples/algorithms/p3o/run_p3o_temperature_0p6_a100x4.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=p3o +export P3O_UPDATE_WEIGHTS_INTERVAL=1 +export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 +export P3O_BEHAVIOR_TEMPERATURE=0.6 +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/examples/algorithms/p3o/run_p3o_temperature_1p2_a100x4.sh b/examples/algorithms/p3o/run_p3o_temperature_1p2_a100x4.sh new file mode 100755 index 000000000..c46b8dc6d --- /dev/null +++ b/examples/algorithms/p3o/run_p3o_temperature_1p2_a100x4.sh @@ -0,0 +1,12 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +export P3O_ALGORITHM=p3o +export P3O_UPDATE_WEIGHTS_INTERVAL=1 +export P3O_ENABLE_TEMPERATURE_OVERRIDE=1 +export P3O_BEHAVIOR_TEMPERATURE=1.2 +source "${SCRIPT_DIR}/common_a100x4.sh" +P3O_run diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..3cac95804 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -49,6 +49,7 @@ ) from relax.utils.distributed_utils import get_gloo_group from relax.utils.env import Envs +from relax.utils.logging_utils import get_logger from relax.utils.memory_utils import clear_memory, print_memory from relax.utils.metrics.metric_utils import compute_rollout_step from relax.utils.opd.opd_utils import ( @@ -93,6 +94,13 @@ from .initialize import init, is_megatron_main_rank from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values from .model import forward_only, initialize_model_and_optimizer, save, train +from .rollout_policy_lag import ( + ROLLOUT_POLICY_TAG, + initial_rollout_policy_snapshot_rollout, + maybe_refresh_rollout_policy, + rollout_weights_tag, + validate_update_weights_interval, +) from .weight_update.common import named_params_and_buffers from .weight_update.train_offload import MegatronTrainStateOffloader from .weight_update.update_weight_from_distributed import UpdateWeightFromDistributed @@ -101,7 +109,7 @@ logging.getLogger("megatron").setLevel(logging.WARNING) -logger = logging.getLogger(__name__) +logger = get_logger(__name__) ROLLOUT_MINI_BATCH_METAS_KEY = "rollout_mini_batch_metas" @@ -248,7 +256,13 @@ def _init( # internally via _switch_model and pushes weights to rollout via # UpdateWeightFromTensor instead of DCS. use_tensor_backuper = not self.args.fully_async or self.args.hybrid + update_weights_interval = validate_update_weights_interval(self.args.update_weights_interval) + if update_weights_interval > 1 and not use_tensor_backuper: + raise ValueError( + "update_weights_interval > 1 requires the synchronous or hybrid TensorBackuper weight-update path" + ) if use_tensor_backuper: + use_rollout_policy_snapshot = update_weights_interval > 1 self.weights_backuper = TensorBackuper.create( source_getter=lambda: named_params_and_buffers( self.args, @@ -256,10 +270,15 @@ def _init( convert_to_global_name=args.megatron_to_hf_mode == "raw", translate_gpu_to_cpu=not self.args.enable_weights_backuper, ), - single_tag=None if args.enable_weights_backuper else "actor", + single_tag=None if args.enable_weights_backuper or use_rollout_policy_snapshot else "actor", ) self._active_model_tag: str | None = "actor" self.weights_backuper.backup("actor") + self._rollout_weights_tag = rollout_weights_tag(update_weights_interval) + # Track the rollout at which rollout policy snapshot was created (for observability) + self._rollout_policy_snapshot_rollout = initial_rollout_policy_snapshot_rollout(start_rollout_id) + if use_rollout_policy_snapshot: + self.weights_backuper.backup(ROLLOUT_POLICY_TAG) if with_ref: self.load_other_checkpoint("ref", args.ref_load) @@ -295,7 +314,7 @@ def _init( self.weight_updater = update_weight_cls( self.args, self.model, - weights_getter=lambda: self.weights_backuper.get("actor"), + weights_getter=lambda: self.weights_backuper.get(self._rollout_weights_tag), model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name, @@ -887,6 +906,8 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # Train if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" + # Store rollout policy snapshot rollout for observability in training metrics + self.args.rollout_policy_snapshot_rollout = self.get_rollout_policy_snapshot_rollout() with timer("actor_train"): train( rollout_id, @@ -969,7 +990,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: if self.args.offload_train: self.sleep() if has_rollout: - self.update_weights() + self.update_weights(rollout_id=rollout_id) tracking_utils.flush_metrics(self.args, compute_rollout_step(self.args, rollout_id)) # RL-only generative eval (uses SGLang via rollout_manager.eval). SFT # uses local eval/predict runner below. @@ -1358,6 +1379,7 @@ def train_hybrid(self, rollout_id) -> None: data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" + self.args.rollout_policy_snapshot_rollout = self.get_rollout_policy_snapshot_rollout() with timer("actor_train"): train( rollout_id, @@ -1440,7 +1462,7 @@ def train_hybrid(self, rollout_id) -> None: self._check_services_health() # Sync weights to rollout via UpdateWeightFromTensor (colocate mode) - self.update_weights() + self.update_weights(rollout_id=rollout_id) tracking_utils.flush_metrics(self.args, compute_rollout_step(self.args, rollout_id)) dist.barrier(group=get_gloo_group()) self._run_step_evaluation(rollout_id, end_update_weight=True) @@ -1607,11 +1629,60 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: if self.args.offload_train and self._per_step_rollout: destroy_process_groups() + def _maybe_refresh_rollout_policy(self, rollout_id: int | None) -> None: + interval = self.args.update_weights_interval + if interval == 1 or rollout_id is None: + return + + if maybe_refresh_rollout_policy( + self.weights_backuper, + rollout_id, + interval, + self.args.num_rollout, + ): + # Store the rollout at which we refreshed the snapshot + self._rollout_policy_snapshot_rollout = rollout_id + 1 + logger.info( + "Refreshed rollout policy snapshot after rollout_id=%s; snapshot version for the next rollout is %s " + "(update_weights_interval=%s)", + rollout_id, + self._rollout_policy_snapshot_rollout, + interval, + ) + else: + next_rollout_snapshot_age = (rollout_id + 1) % interval + logger.info( + "Retaining rollout policy snapshot after rollout_id=%s; next rollout snapshot age will be %s " + "rollout(s) " + "(update_weights_interval=%s)", + rollout_id, + next_rollout_snapshot_age, + interval, + ) + + def get_rollout_policy_snapshot_rollout(self) -> int: + """Return the rollout at which the current rollout policy snapshot was + created. + + Returns 0 for on-policy (interval=1) or when snapshot tracking is + unavailable. + """ + return getattr(self, "_rollout_policy_snapshot_rollout", 0) + @timer - def update_weights(self) -> None: + def update_weights(self, rollout_id: int | None = None) -> None: + """Publish the selected actor snapshot to rollout workers. + + Args: + rollout_id: Zero-based rollout identifier that controls periodic + rollout-policy snapshot refreshes. ``None`` skips refresh + bookkeeping for callers outside the rollout loop. + """ if self.args.debug_train_only or self.args.debug_rollout_only: return + self._maybe_refresh_rollout_policy(rollout_id) + if self.args.offload_train: # CRITICAL: Barrier before onload_weights to ensure ALL ranks have # completed sleep() (and released GPU memory via tms.pause()) before diff --git a/relax/backends/megatron/cp_utils.py b/relax/backends/megatron/cp_utils.py index 98129f001..542994d16 100644 --- a/relax/backends/megatron/cp_utils.py +++ b/relax/backends/megatron/cp_utils.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from collections.abc import Callable, Sequence import torch import torch.distributed as dist @@ -13,6 +15,14 @@ mpu = None +def _validate_metadata_lengths(**metadata: Sequence[object] | None) -> None: + """Reject CP metadata lists that would otherwise be silently truncated.""" + lengths = {name: len(values) for name, values in metadata.items() if values is not None} + if len(set(lengths.values())) > 1: + formatted = ", ".join(f"{name}={length}" for name, length in lengths.items()) + raise ValueError(f"CP metadata lengths must match; got {formatted}") + + def maybe_padded_total_lengths( total_lengths: list[int], qkv_format: str, @@ -48,19 +58,20 @@ def get_logits_and_tokens_offset_with_cp( """All offsets start from the begining of the prompt.""" cp_rank = dynamic_cp_rank if dynamic_cp_rank is not None else mpu.get_context_parallel_rank() cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() - assert cp_size > 1 + if cp_size <= 1: + raise ValueError(f"Context parallel size must be > 1, got {cp_size}") prompt_length = total_length - response_length if padded_total_length is not None: # Bridge VL+CP+thd: per-sample padded length is already aligned to tp*cp*2. - assert padded_total_length % (2 * cp_size) == 0, ( - f"padded_total_length={padded_total_length} not divisible by 2*cp={2 * cp_size}" - ) + if padded_total_length % (2 * cp_size) != 0: + raise ValueError(f"padded_total_length={padded_total_length} not divisible by 2*cp={2 * cp_size}") chunk_size = padded_total_length // (2 * cp_size) elif qkv_format == "thd": chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size) else: - assert max_seq_len is not None, "max_seq_len must be provided for qkv_format=bshd" + if max_seq_len is None: + raise ValueError("max_seq_len must be provided for qkv_format=bshd") chunk_size = (max_seq_len + 2 * cp_size - 1) // (2 * cp_size) # the offset of 2 chunks @@ -99,6 +110,13 @@ def get_sum_of_sample_mean( dynamic_cp_rank: int | None = None, ) -> Callable[[torch.Tensor], torch.Tensor]: """Calculate correct sample mean for CP.""" + _validate_metadata_lengths( + total_lengths=total_lengths, + response_lengths=response_lengths, + loss_masks=loss_masks, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, + ) cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() if cp_size == 1: @@ -106,7 +124,7 @@ def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: return sum( [ (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1) - for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False) + for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=True) ] ) @@ -114,7 +132,7 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: return sum( [ (x_i * loss_mask_i).sum() - for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False) + for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=True) ] ) @@ -123,7 +141,7 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: chunked_loss_masks: list[torch.Tensor] = [] for i, (total_length, response_length, loss_mask) in enumerate( - zip(total_lengths, response_lengths, loss_masks, strict=False) + zip(total_lengths, response_lengths, loss_masks, strict=True) ): max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None @@ -147,7 +165,7 @@ def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: [ (x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1) for x_i, chunked_loss_mask, loss_mask in zip( - x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks, strict=False + x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks, strict=True ) ] ) @@ -157,7 +175,7 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: [ (x_i * chunked_loss_mask).sum() for x_i, chunked_loss_mask in zip( - x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, strict=False + x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, strict=True ) ] ) @@ -192,6 +210,13 @@ def get_cp_local_num_tokens( For ``cp_size == 1`` this reduces to the total number of unmasked tokens (preserving the historical per-sample ``clamp_min(., 1)``). """ + _validate_metadata_lengths( + total_lengths=total_lengths, + response_lengths=response_lengths, + loss_masks=loss_masks, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, + ) cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() if cp_size == 1: return sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in loss_masks]) @@ -200,7 +225,7 @@ def get_cp_local_num_tokens( # counted tokens exactly match the ones sum_of_token contributes on this rank. total: torch.Tensor | None = None for i, (total_length, response_length, loss_mask) in enumerate( - zip(total_lengths, response_lengths, loss_masks, strict=False) + zip(total_lengths, response_lengths, loss_masks, strict=True) ): max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None @@ -225,6 +250,67 @@ def get_cp_local_num_tokens( return total +def get_cp_local_valid_mask( + total_lengths: list[int], + response_lengths: list[int], + loss_masks: list[torch.Tensor], + qkv_format: str = "thd", + max_seq_lens: list[int] | None = None, + padded_total_lengths: list[int] | None = None, + dynamic_cp_size: int | None = None, + dynamic_cp_rank: int | None = None, +) -> torch.Tensor: + """Build the CP-local boolean mask of loss-contributing response tokens. + + Returns a single 1-D mask over this rank's concatenated response tokens, + aligned with the layout that ``get_sum_of_sample_mean`` reduces over. Callers + that must compute a statistic and a loss over *identical* token sets (P3O's + ESS pre-pass and its loss) share this helper instead of re-deriving the + zig-zag slicing, which is where the two can silently drift apart. + + For ``cp_size == 1`` this is just the concatenation of ``loss_masks``. + """ + _validate_metadata_lengths( + total_lengths=total_lengths, + response_lengths=response_lengths, + loss_masks=loss_masks, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, + ) + cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() + if cp_size == 1: + return torch.cat([loss_mask.bool() for loss_mask in loss_masks], dim=0) + + chunks: list[torch.Tensor] = [] + for i, (total_length, response_length, loss_mask) in enumerate( + zip(total_lengths, response_lengths, loss_masks, strict=True) + ): + max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None + padded_total_length = padded_total_lengths[i] if padded_total_lengths is not None else None + prompt_length = total_length - response_length + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp( + total_length, + response_length, + qkv_format, + max_seq_len, + padded_total_length, + dynamic_cp_size=dynamic_cp_size, + dynamic_cp_rank=dynamic_cp_rank, + ) + loss_mask_0 = loss_mask[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + loss_mask_1 = loss_mask[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + chunks.append(torch.cat([loss_mask_0, loss_mask_1], dim=0).bool()) + + if not chunks: + if not loss_masks: + raise ValueError( + "P3O cp_utils: both loss_masks and computed chunks are empty; " + "cannot determine device for the returned tensor." + ) + return torch.zeros(0, dtype=torch.bool, device=loss_masks[0].device) + return torch.cat(chunks, dim=0) + + def all_gather_with_cp( tensor: torch.Tensor, total_length: int, @@ -260,7 +346,9 @@ def all_gather_with_cp( chunk_0 = tensor[: logits_offset[0][1] - logits_offset[0][0]] chunk_1 = tensor[logits_offset[0][1] - logits_offset[0][0] :] - assert chunk_1.shape[0] == logits_offset[1][1] - logits_offset[1][0] + expected_chunk_1_len = logits_offset[1][1] - logits_offset[1][0] + if chunk_1.shape[0] != expected_chunk_1_len: + raise ValueError(f"chunk_1 length {chunk_1.shape[0]} != expected {expected_chunk_1_len}") def zero(len: int) -> torch.Tensor: return torch.zeros( @@ -290,7 +378,8 @@ def zero(len: int) -> torch.Tensor: right = zero(total_length - 1 - logits_offset[1][1]) full_tensor = torch.cat([left, chunk_0, mid, chunk_1, right], dim=0) - assert full_tensor.shape[0] == response_length, f"Expected {response_length}, got {full_tensor.shape}" + if full_tensor.shape[0] != response_length: + raise ValueError(f"Expected response_length={response_length}, got shape {full_tensor.shape}") full_tensor = dist.nn.all_reduce(full_tensor, group=cp_group) return full_tensor @@ -307,7 +396,8 @@ def slice_with_cp( cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() if qkv_format == "bshd": - assert max_seq_len is not None + if max_seq_len is None: + raise ValueError("max_seq_len is required when qkv_format=bshd") def pad_tokens(tokens, pad): if isinstance(pad_value, Callable): @@ -351,10 +441,11 @@ def slice_log_prob_with_cp( dynamic_cp_size: int | None = None, dynamic_cp_rank: int | None = None, ) -> list[float] | torch.Tensor: - assert len(log_prob) == response_length, ( - f"log_prob length mismatch: len(log_prob)={len(log_prob)}, " - f"response_length={response_length}, total_length={total_length}" - ) + if len(log_prob) != response_length: + raise ValueError( + f"log_prob length mismatch: len(log_prob)={len(log_prob)}, " + f"response_length={response_length}, total_length={total_length}" + ) cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() @@ -477,7 +568,8 @@ def _nccl_all_gather_variable_tensors( Every rank in ``group`` must call this with a non-empty ``values`` so the collective is symmetric and a device/dtype is available. """ - assert values, "_nccl_all_gather_variable_tensors requires a non-empty values list on every rank" + if not values: + raise ValueError("_nccl_all_gather_variable_tensors requires a non-empty values list on every rank") local_sizes = torch.tensor([v.shape[0] for v in values], dtype=torch.long, device=values[0].device) num_samples = torch.tensor([len(values)], dtype=torch.long, device=values[0].device) @@ -558,6 +650,12 @@ def dynamic_cp_merge_output( if dynamic_cp_size > 1: dynamic_cp_group = mpu.get_dynamic_data_context_parallel_groups(group_size=dynamic_cp_size) ptls = padded_total_lengths if padded_total_lengths is not None else [None] * len(values) + _validate_metadata_lengths( + values=values, + total_lengths=total_lengths, + response_lengths=response_lengths, + padded_total_lengths=ptls, + ) values = [ all_gather_with_cp( v, @@ -568,7 +666,7 @@ def dynamic_cp_merge_output( dynamic_cp_rank=dynamic_cp_rank, dynamic_cp_group=dynamic_cp_group, ) - for v, tl, rl, ptl in zip(values, total_lengths, response_lengths, ptls, strict=False) + for v, tl, rl, ptl in zip(values, total_lengths, response_lengths, ptls, strict=True) ] # 2. collect all sub-groups' samples across the static CP group and reorder. @@ -581,10 +679,11 @@ def dynamic_cp_merge_output( # A subdivided mb always carries a partition order; reorder back to the # original mb sample order so the write-back aligns with micro_batch_indices. # Fail loud (not a silent wrong order) if the invariant ever breaks. - assert partition_order is not None and len(partition_order) == len(values), ( - "dynamic-CP merge: partition_order missing or length mismatch " - f"(order={None if partition_order is None else len(partition_order)}, values={len(values)})" - ) + if partition_order is None or len(partition_order) != len(values): + raise ValueError( + "dynamic-CP merge: partition_order missing or length mismatch " + f"(order={None if partition_order is None else len(partition_order)}, values={len(values)})" + ) reordered: list = [None] * len(values) for new_pos, orig_pos in enumerate(partition_order): reordered[orig_pos] = values[new_pos] diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index 9aff3da6f..374c3d951 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -250,6 +250,46 @@ def pad_and_flatten( return torch.cat(padded_list, dim=0), num_items +def build_rl_forward_kwargs(args: Namespace, batch: dict[str, Any]) -> tuple[dict[str, Any], bool]: + """Build policy-forward inputs shared by training and the P3O stats pass. + + The returned flag identifies layouts that require unsplit tokens, so + callers can select the matching dynamic-CP process group around their own + forward. Keeping that process-group mutation outside this pure helper + preserves each caller's distinct lifetime while keeping the model inputs + identical. + """ + is_vl_model = bool(getattr(args, "is_vl_model", False)) + multimodal_inputs = batch.get("multimodal_train_inputs") + needs_unsplit = is_vl_model or multimodal_inputs is not None or getattr(args, "uses_unsplit_forward", False) + use_unsplit = needs_unsplit and "unsplit_tokens" in batch + + forward_kwargs = { + "input_ids": batch["unsplit_tokens"] if use_unsplit else batch["tokens"], + "position_ids": None, + "attention_mask": None, + "labels": None, + "packed_seq_params": None if use_unsplit else batch["packed_seq_params"], + "loss_mask": batch["full_loss_masks"], + } + + # THD VL+CP uses the bridge-specific attention mask and packed layout. The + # external RL loss consumes full_loss_masks, so GPTModel must not compute an + # internal loss for this bridge path. + if needs_unsplit and "vlm_packed_seq_params" in batch: + forward_kwargs["attention_mask"] = batch["unsplit_attention_mask"] + forward_kwargs["packed_seq_params"] = batch["vlm_packed_seq_params"] + forward_kwargs["loss_mask"] = None + + # Text-only batches for a VL checkpoint deliberately carry no multimodal + # kwargs. Custom multimodal data with is_vl_model=False follows the same + # gate in both the stats and gradient passes. + if is_vl_model and multimodal_inputs: + forward_kwargs.update(multimodal_inputs) + + return forward_kwargs, needs_unsplit + + def get_batch( data_iterator: "DataIterator", keys: Sequence[str], @@ -702,6 +742,24 @@ def reset(self) -> "DataIterator": self.offset = 0 return self + def snapshot_position(self) -> int: + """Return the current offset so it can be restored later. + + ``reset()`` rewinds to the start of the whole rollout, which is wrong + for replaying a single optimizer window that begins mid-rollout. P3O's + ESS pre-pass consumes the window once and must hand the iterator back + exactly where it found it. + """ + return self.offset + + def restore_position(self, position: int) -> None: + """Restore an offset previously returned by :meth:`snapshot_position`. + + Works for both the fixed micro-batch-size and the explicit + ``micro_batch_indices`` schedule, including non-zero start offsets. + """ + self.offset = position + def get_data_iterator( args: Namespace, diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 372a139cf..52effa8ba 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -1,3 +1,5 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + from argparse import Namespace from collections.abc import Callable, Iterator from functools import partial @@ -17,8 +19,16 @@ compute_policy_opd_loss, resolve_opd_gather_topk_token_ids, validate_opd_topk_gather, + validate_p3o_opd_compatibility, +) +from relax.utils.training.p3o_utils import ( + P3OStepContext, + compute_p3o_sufficient_stats_unchecked, + compute_p3o_token_terms, + finalize_p3o_step_context, ) from relax.utils.training.ppo_utils import ( + GRPO_STYLE_ADVANTAGE_ESTIMATORS, calculate_log_probs_and_entropy, compute_approx_kl, compute_cispo_loss, @@ -37,11 +47,13 @@ from .cp_utils import ( all_gather_with_cp, get_cp_local_num_tokens, + get_cp_local_valid_mask, get_logits_and_tokens_offset_with_cp, get_sum_of_sample_mean, maybe_padded_total_lengths, slice_log_prob_with_cp, ) +from .p3o_step import synchronize_p3o_stats def get_responses( @@ -575,7 +587,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) for i in range(len(log_probs)) ] - if args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo"]: + if args.advantage_estimator in GRPO_STYLE_ADVANTAGE_ESTIMATORS: rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) returns = get_grpo_returns(rewards, kl) # TODO: is the copy necessary? @@ -792,6 +804,208 @@ def icepop_function( return pg_loss, loss_masks, metrics +def get_p3o_step_context(args: Namespace) -> P3OStepContext: + """Fetch the frozen P3O context for the optimizer step in progress. + + The context is published by the Megatron backend's ESS pre-pass + (``model.py::compute_p3o_step_context``) before the training + forward/backward schedule starts, and is deliberately not passed through + the micro-batch dict: every micro-batch of the step must see the exact same + cap. + """ + step_context = getattr(args, "_p3o_step_context", None) + if step_context is None: + raise RuntimeError( + "P3O: no optimizer-step context available. The ESS pre-pass must run " + "before the training forward/backward schedule." + ) + return step_context + + +def get_p3o_context( + args: Namespace, + log_probs: torch.Tensor, + behavior_log_probs: torch.Tensor, + valid_mask: torch.Tensor, +) -> P3OStepContext: + """Resolve the configured P3O ESS scope for one loss micro-batch.""" + scope = getattr(args, "p3o_ess_scope", "micro-batch") + if scope == "step": + return get_p3o_step_context(args) + if scope != "micro-batch": + raise ValueError(f"P3O ESS scope must be 'micro-batch' or 'step', got {scope!r}") + + stats, invalid_count = compute_p3o_sufficient_stats_unchecked( + log_probs, + behavior_log_probs, + valid_mask, + ) + distributed = dist.is_available() and dist.is_initialized() + stats = synchronize_p3o_stats( + stats, + invalid_count, + dp_cp_group=mpu.get_data_parallel_group(with_context_parallel=True) if distributed else None, + pp_group=None, + is_pipeline_last_stage=True, + ) + return finalize_p3o_step_context(stats) + + +def p3o_loss_function( + args: Namespace, + batch: RolloutBatch, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute the P3O loss and metrics for one micro-batch. + + P3O is kept out of :func:`policy_loss_function` on purpose. Its objective is + a score-function update whose ratio coefficient is fully detached and capped + by the optimizer-step ESS, so none of the PPO machinery applies: no + advantage-sign branch, no lower clip bound, and ``eps_clip`` has no effect. + Mixing it into the PPO branch would mean threading a "which clipping regime" + flag through code that assumes a two-sided surrogate. + + The behavior policy is the rollout sampling distribution + (``rollout_log_probs``), never a detached copy of the current forward: + substituting the latter would erase exactly the policy lag / temperature + mismatch P3O exists to absorb. + + Args: + args: Configuration. Reads ``entropy_coef``, ``use_kl_loss`` / + ``kl_loss_coef`` (frozen-reference regularization, reported + separately from the adaptive behavior KL), and the P3O step context. + batch: Mini-batch with "advantages", "rollout_log_probs", + "unconcat_tokens", "total_lengths", "response_lengths", "loss_masks". + logits: Policy logits with shape ``[1, T, V]``. + sum_of_sample_mean: Reduction over this micro-batch's tokens. P3O + requires the token-sum variant (``--calculate-per-token-loss``) so + that per-micro-batch denominators do not re-enter the objective. + + Returns: + Tuple of ``(loss, metrics)``. Metric keys are prefixed ``p3o/`` except + the shared ``loss`` / ``pg_loss`` / ``entropy_loss`` keys kept for + dashboard compatibility. Global scalars (ESS, cap, ratio moments) are + pre-multiplied by this rank's valid-token count, because the caller + divides every reported metric by the globally reduced token count. + """ + if isinstance(batch["advantages"], list): + advantages = torch.cat(batch["advantages"], dim=0) + else: + advantages = batch["advantages"] + + # Raise, not assert: under `python -O` a stripped check would fall through to + # a KeyError deep in the loss, or worse, a silently wrong behavior policy. + if batch.get("rollout_log_probs") is None: + raise ValueError( + "P3O requires actual rollout log-probs as the behavior policy; run with --use-rollout-logprobs." + ) + + total_lengths = batch["total_lengths"] + response_lengths = batch["response_lengths"] + max_seq_lens = batch.get("max_seq_lens", None) + padded_total_lengths = batch.get("padded_total_lengths", None) + + _, log_probs_and_entropy = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=True, + max_seq_lens=max_seq_lens, + padded_total_lengths=padded_total_lengths, + dynamic_cp_size=batch.get("dynamic_cp_size", None), + dynamic_cp_rank=batch.get("dynamic_cp_rank", None), + ) + + log_probs = torch.cat(log_probs_and_entropy["log_probs"], dim=0) + behavior_log_probs = torch.cat(batch["rollout_log_probs"], dim=0) + + valid_mask = get_cp_local_valid_mask( + total_lengths, + response_lengths, + batch["loss_masks"], + args.qkv_format, + max_seq_lens, + padded_total_lengths, + dynamic_cp_size=batch.get("dynamic_cp_size", None), + dynamic_cp_rank=batch.get("dynamic_cp_rank", None), + ) + step_context = get_p3o_context(args, log_probs, behavior_log_probs, valid_mask) + + terms = compute_p3o_token_terms( + log_probs=log_probs, + behavior_log_probs=behavior_log_probs, + advantages=advantages, + valid_mask=valid_mask, + step_context=step_context, + kl_mode=getattr(args, "p3o_kl_mode", "proxy"), + clip_low=getattr(args, "clip_low", 0.2), + clip_high=getattr(args, "clip_high", 0.2), + ) + + score_loss = sum_of_sample_mean(terms.score_loss) + adaptive_kl_loss = sum_of_sample_mean(terms.adaptive_kl_loss) + # behavior_kl_proxy: sampled-token k3 proxy (1-ESS), not full-vocabulary KL. + # Measures concentration of importance ratios via ESS, not distributional shift. + behavior_kl_proxy = sum_of_sample_mean(terms.behavior_kl_proxy) + # cap_fraction: fraction of tokens where adaptive cap binds (ratio > ESS). + # Different from PPO's clip_fraction which measures fixed-interval clipping. + cap_fraction = sum_of_sample_mean(terms.cap_hits) + clip_fraction = sum_of_sample_mean(terms.clip_hits) + + entropy = torch.cat(log_probs_and_entropy["entropy"], dim=0) + entropy_loss = sum_of_sample_mean(entropy) + + loss = score_loss + adaptive_kl_loss - args.entropy_coef * entropy_loss + + reference_kl_loss = None + reference_kl_metric = loss.detach().new_zeros(()) + if args.use_kl_loss: + # Optional frozen-reference regularization. Orthogonal to the adaptive + # behavior KL above and reported under its own key. + ref_log_probs = torch.cat(batch["ref_log_probs"], dim=0) + reference_kl = compute_approx_kl(log_probs, ref_log_probs, kl_loss_type=args.kl_loss_type) + reference_kl_loss = sum_of_sample_mean(reference_kl) + reference_kl_metric = reference_kl_loss.clone().detach() + loss = loss + args.kl_loss_coef * reference_kl_loss + + if log_probs.numel() == 0: + loss += 0 * logits.sum() + + # Global step scalars are reported as scalar * local_valid_tokens so that the + # caller's divide-by-global-token-count recovers the scalar itself. + local_valid_tokens = valid_mask.sum().to(torch.float32) + + def scaled(value: torch.Tensor) -> torch.Tensor: + return (value.to(torch.float32) * local_valid_tokens).clone().detach() + + reported_loss = { + "loss": loss.clone().detach(), + "pg_loss": score_loss.clone().detach(), + "entropy_loss": entropy_loss.clone().detach(), + "p3o/score_loss": score_loss.clone().detach(), + "p3o/behavior_kl_proxy": behavior_kl_proxy.clone().detach(), + "p3o/adaptive_kl_loss": adaptive_kl_loss.clone().detach(), + "p3o/reference_kl": reference_kl_metric, + "p3o/entropy": entropy_loss.clone().detach(), + "p3o/cap_fraction": cap_fraction.clone().detach(), + "p3o/clip_fraction": clip_fraction.clone().detach(), + "p3o/total_loss": loss.clone().detach(), + "p3o/normalized_ess": scaled(step_context.normalized_ess), + "p3o/adaptive_cap": scaled(step_context.adaptive_cap), + "p3o/ratio_mean": scaled(step_context.ratio_mean), + "p3o/ratio_std": scaled(step_context.ratio_std), + "p3o/valid_tokens": scaled(step_context.valid_token_count), + } + + if reference_kl_loss is not None: + reported_loss["kl_loss"] = reference_kl_loss.clone().detach() + + return loss, reported_loss + + def _get_reinforce_plus_plus_mask_safe_reducer( reducer: Callable[[torch.Tensor], torch.Tensor], loss_masks: list[torch.Tensor], @@ -1304,6 +1518,24 @@ def sft_loss_function_chunked( return loss, {"loss": loss.clone().detach()} +def _select_policy_loss_function( + args: Namespace, +) -> Callable[..., tuple[torch.Tensor, dict[str, torch.Tensor]]]: + """Select one policy objective without composing unrelated algorithm + families. + + P3O has a dedicated score-function/trust-region objective and therefore + bypasses :func:`policy_loss_function`, including its optional + :func:`compute_policy_opd_loss` term. The compatibility guard is repeated + here so callers that bypass normal argument validation still fail before a + hybrid loss can be computed. + """ + validate_p3o_opd_compatibility(args) + if getattr(args, "advantage_estimator", None) == "p3o": + return p3o_loss_function + return policy_loss_function + + def loss_function( args: Namespace, batch: RolloutBatch, @@ -1347,16 +1579,26 @@ def loss_function( # normalizer is correct even when CP differs across micro-batches (dynamic CP). # Under static CP it equals the old full-sample count distributed across ranks, # so the final loss/grad/metric are unchanged after all-reduce. - num_tokens = get_cp_local_num_tokens( + token_count_args = ( batch["total_lengths"], batch["response_lengths"], batch["loss_masks"], args.qkv_format, batch.get("max_seq_lens", None), batch.get("padded_total_lengths", None), - dynamic_cp_size=batch.get("dynamic_cp_size", None), - dynamic_cp_rank=batch.get("dynamic_cp_rank", None), ) + token_count_kwargs = { + "dynamic_cp_size": batch.get("dynamic_cp_size", None), + "dynamic_cp_rank": batch.get("dynamic_cp_rank", None), + } + if getattr(args, "advantage_estimator", None) == "p3o": + # P3O's optimizer-step objective is normalized by the exact global count + # used for ESS. The generic helper preserves a historical clamp-to-one + # for fully masked samples when CP=1, which would create phantom tokens + # and make the final loss depend on the CP partition. + num_tokens = get_cp_local_valid_mask(*token_count_args, **token_count_kwargs).sum() + else: + num_tokens = get_cp_local_num_tokens(*token_count_args, **token_count_kwargs) num_samples = len(batch["response_lengths"]) sum_of_sample_mean = get_sum_of_sample_mean( @@ -1373,7 +1615,7 @@ def loss_function( match args.loss_type: case "policy_loss": - func = policy_loss_function + func = _select_policy_loss_function(args) case "value_loss": func = value_loss_function case "sft": diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 077e461d9..c5bfc26ff 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. +import contextlib import dataclasses import gc import math @@ -45,9 +46,10 @@ ) from .checkpoint import load_checkpoint, save_checkpoint -from .data import DataIterator, get_batch +from .data import DataIterator, build_rl_forward_kwargs, get_batch from .loss import loss_function from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze +from .rollout_policy_lag import build_rollout_policy_age_metrics logger = get_logger(__name__) @@ -171,6 +173,23 @@ def _chunked_call(input_, weight=None, runtime_gather_output=None): output_layer.forward = original_forward +@contextmanager +def _preserved_dynamic_cp_group(args: Namespace, model: Sequence[torch.nn.Module]) -> Iterator[None]: + """Restore the static context-parallel group after dynamic-CP forwards.""" + if not getattr(args, "dynamic_context_parallel", False): + yield + return + + inner = model[0] + while hasattr(inner, "module"): + inner = inner.module + original_cp_group = inner.pg_collection.cp + try: + yield + finally: + inner.pg_collection.cp = original_cp_group + + def _should_use_sft_chunked(args: Namespace) -> bool: """Gate for the SFT chunked-logits path. @@ -1062,36 +1081,9 @@ def forward_step( loss_mask=batch["full_loss_masks"], ) else: - has_mm_inputs = batch.get("multimodal_train_inputs", None) is not None - needs_unsplit = is_vl_model or has_mm_inputs or getattr(args, "uses_unsplit_forward", False) - use_unsplit = needs_unsplit and "unsplit_tokens" in batch - - forward_kwargs = { - "input_ids": batch["unsplit_tokens"] if use_unsplit else batch["tokens"], - "position_ids": None, - "attention_mask": None, - "labels": None, - "packed_seq_params": None if use_unsplit else batch["packed_seq_params"], - "loss_mask": batch["full_loss_masks"], - } - - # thd VL+CP: bridge needs per-sample attention_mask + matching thd - # packed_seq_params (align_size = tp*cp*2). loss_mask is None - # because labels=None means GPTModel won't run internal loss; - # Relax's loss is computed externally from full_loss_masks. - if needs_unsplit and "vlm_packed_seq_params" in batch: - forward_kwargs["attention_mask"] = batch["unsplit_attention_mask"] - forward_kwargs["packed_seq_params"] = batch["vlm_packed_seq_params"] - forward_kwargs["loss_mask"] = None - + forward_kwargs, needs_unsplit = build_rl_forward_kwargs(args, batch) _attach_mtp_forward_kwargs(args, batch, forward_kwargs) - # VL model with text-only batch has is_vl_model=True but no - # multimodal_train_inputs in batch — no kwargs to splice in. - mm_inputs = batch.get("multimodal_train_inputs") - if is_vl_model and mm_inputs: - forward_kwargs.update(mm_inputs) - # Dynamic CP: point pg_collection.cp at this mb's dynamic CP sub-group for # the VL bridge forward. Set every mb (incl. size 1) to avoid a stale group # leaking from a previous mb; restored once after forward_backward (below). @@ -1126,15 +1118,6 @@ def forward_step( # and lm_head_forward are set. return output_tensor, partial(loss_function, args, batch, num_microbatches, lm_head_forward=lm_head_forward) - # Dynamic CP: forward_step overwrites pg_collection.cp per micro-batch (VL bridge); - # save the original static CP group here and restore after forward+backward. - _dcp_orig_cp_group = None - if getattr(args, "dynamic_context_parallel", False): - inner = model[0] - while hasattr(inner, "module"): - inner = inner.module - _dcp_orig_cp_group = inner.pg_collection.cp - # Forward pass. use_streaming = ( getattr(args, "use_dynamic_batch_size", False) @@ -1158,19 +1141,41 @@ def forward_step( forward_backward_func = streaming_forward_backward_pipelining_without_interleaving else: forward_backward_func = get_forward_backward_func() - losses_reduced = forward_backward_func( - forward_step_func=forward_step, - data_iterator=data_iterator, - model=model, - num_microbatches=num_microbatches, - seq_length=args.seq_length, - micro_batch_size=args.micro_batch_size, - decoder_seq_length=args.decoder_seq_length, - forward_only=False, - ) - if _dcp_orig_cp_group is not None: - inner.pg_collection.cp = _dcp_orig_cp_group + # Dynamic CP mutates the model's CP process group inside each forward. + # Protect both P3O passes so failures cannot leak a per-micro-batch group. + with _preserved_dynamic_cp_group(args, model): + # Optional step scope freezes one adaptive cap before gradients are + # produced. Micro-batch scope computes its cap inside the loss callback. + p3o_context_manager = contextlib.nullcontext() + if ( + getattr(args, "advantage_estimator", None) == "p3o" + and getattr(args, "p3o_ess_scope", "micro-batch") == "step" + ): + from relax.backends.megatron.p3o_step import ( + compute_p3o_step_context, + p3o_step_context_published, + ) + + p3o_step_context = compute_p3o_step_context( + args=args, + data_iterator=data_iterator, + model=model, + num_microbatches=num_microbatches, + ) + p3o_context_manager = p3o_step_context_published(args, p3o_step_context) + + with p3o_context_manager: + losses_reduced = forward_backward_func( + forward_step_func=forward_step, + data_iterator=data_iterator, + model=model, + num_microbatches=num_microbatches, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + decoder_seq_length=args.decoder_seq_length, + forward_only=False, + ) # CI check: verify only MTP parameters have non-zero gradients when truncation happens # This check must happen before optimizer.step() as gradients may be modified during step @@ -1457,6 +1462,18 @@ def train( log_dict[f"train/{role_tag}cur_epoch"] = (accumulated_step_id + 1) / ( num_per_epoch * num_steps_per_rollout ) + + # P3O observability: track rollout policy age + if getattr(args, "advantage_estimator", None) == "p3o" and args.update_weights_interval > 1: + snapshot_rollout = getattr(args, "rollout_policy_snapshot_rollout", 0) + current_rollout = rollout_id + log_dict.update( + build_rollout_policy_age_metrics( + current_rollout_id=current_rollout, + rollout_policy_snapshot_rollout=snapshot_rollout, + ) + ) + tracking_utils.log(args, log_dict, step_key="train/step") tracking_utils.flush_metrics(args, accumulated_step_id) diff --git a/relax/backends/megatron/p3o_step.py b/relax/backends/megatron/p3o_step.py new file mode 100644 index 000000000..fbf3e677a --- /dev/null +++ b/relax/backends/megatron/p3o_step.py @@ -0,0 +1,289 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Optional optimizer-step scoped ESS pre-pass for P3O. + +Relax computes ESS over one whole optimizer step to ensure that neither the +number of micro-batches nor the DP/CP split change the adaptive cap or the +final loss. 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. Relax's approach provides partition invariance: + + stats pass (no grad) over every micro-batch of the window + -> local S1 / S2 / N + -> one all-reduce over DP x CP + -> immutable P3OStepContext + train pass over the same data, same RNG, one frozen cap + -> token-sum loss, global-token normalization + +The pre-pass replays the same iterator window, so it snapshots and restores both +the iterator offsets and the RNG state. Anything that mutates state during a +no-grad forward (dropout, FP8 amax history) would break that replay and is +rejected in ``arguments.py`` rather than silently tolerated here. +""" + +from argparse import Namespace +from collections.abc import Iterator, Sequence +from contextlib import contextmanager + +import torch +from megatron.core import mpu +from megatron.core.pipeline_parallel import get_forward_backward_func + +from relax.utils.logging_utils import get_logger +from relax.utils.training.p3o_replay import preserved_iterator_positions, preserved_rng_state +from relax.utils.training.p3o_utils import ( + P3OStepContext, + P3OSufficientStats, + compute_p3o_sufficient_stats_unchecked, + finalize_p3o_step_context, +) + +from .cp_utils import get_cp_local_valid_mask +from .data import DataIterator, build_rl_forward_kwargs, get_batch + + +logger = get_logger(__name__) + +P3O_STEP_CONTEXT_ATTR = "_p3o_step_context" +P3O_NONFINITE_RATIO_ERROR = ( + "P3O: non-finite importance ratio at a valid response token on at least one rank; " + "refusing to silently fall back to ESS=1. Check rollout log-probs and mask alignment." +) + + +def _local_stats_from_batch( + args: Namespace, batch: dict, log_probs: list[torch.Tensor] +) -> tuple[P3OSufficientStats, torch.Tensor]: + """Accumulate one micro-batch's ESS contribution from its log-probs. + + Returns: + ``(stats, invalid_flag)``, where ``invalid_flag`` is a device-resident + ``float64`` scalar set to ``1.0`` if this micro-batch produced a + non-finite ratio. It is reduced with ``S1/S2/N`` rather than checked + here, so the pre-pass adds no GPU-CPU sync per micro-batch. + """ + if batch.get("__is_dummy__", False): + # Dummy micro-batches exist only to align num_microbatches across DP + # ranks; they must contribute nothing to S1 / S2 / N. + device = log_probs[0].device if log_probs else "cpu" + return ( + P3OSufficientStats.zeros(device=device), + torch.zeros((), dtype=torch.float64, device=device), + ) + + total_lengths = batch["total_lengths"] + response_lengths = batch["response_lengths"] + padded_total_lengths = batch.get("padded_total_lengths", None) + + current = torch.cat(log_probs, dim=0) + behavior = torch.cat(batch["rollout_log_probs"], dim=0) + valid_mask = get_cp_local_valid_mask( + total_lengths, + response_lengths, + batch["loss_masks"], + args.qkv_format, + batch.get("max_seq_lens", None), + padded_total_lengths, + dynamic_cp_size=batch.get("dynamic_cp_size", None), + dynamic_cp_rank=batch.get("dynamic_cp_rank", None), + ) + return compute_p3o_sufficient_stats_unchecked(current, behavior, valid_mask) + + +def synchronize_p3o_stats( + stats: P3OSufficientStats, + invalid_count: torch.Tensor, + *, + dp_cp_group: torch.distributed.ProcessGroup | None, + pp_group: torch.distributed.ProcessGroup | None, + is_pipeline_last_stage: bool, +) -> P3OSufficientStats: + """Reduce last-stage stats over DP x CP, then publish them over PP. + + Pipeline-last is the only stage with logits. It first sums ``S1/S2/N`` and + the invalid-ratio flag over DP x CP. The already-global vector is then + broadcast, never summed, over PP so every stage finalizes the same context. + TP replicas use independent but equivalent groups. Process groups are + supplied by the caller so the collective scope is explicit at the runtime + integration boundary. + """ + vector = torch.cat((stats.as_vector(), invalid_count.reshape(1).to(dtype=torch.float64))) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + if is_pipeline_last_stage: + torch.distributed.all_reduce(vector, op=torch.distributed.ReduceOp.SUM, group=dp_cp_group) + + if pp_group is not None: + torch.distributed.broadcast( + vector, + group=pp_group, + group_src=torch.distributed.get_world_size(group=pp_group) - 1, + ) + + valid = vector[3] <= 0 + if valid.device.type == "cpu": + if not bool(valid): + raise ValueError(P3O_NONFINITE_RATIO_ERROR) + else: + # Keep the accelerator hot path asynchronous. Every rank observes the + # globally reduced invalid flag, so they all fail consistently. + torch._assert_async(valid, P3O_NONFINITE_RATIO_ERROR) + return P3OSufficientStats.from_vector(vector[:3]) + + +def compute_p3o_step_context( + args: Namespace, + data_iterator: Sequence[DataIterator], + model: Sequence[torch.nn.Module], + num_microbatches: int, +) -> P3OStepContext: + """Run the no-grad stats pass and return this step's frozen P3O context. + + Args: + args: Runtime arguments. + data_iterator: The same iterator(s) the training pass will consume. + model: DDP-wrapped model chunks. + num_microbatches: Micro-batch count for this optimizer step. + + Returns: + The immutable :class:`P3OStepContext` for the step. + """ + from .loss import get_log_probs_and_entropy + + # Accumulated in a cell rather than a rebound local: the write happens inside + # the nested loss callback that Megatron's schedule invokes, one level deeper + # than forward_step. + stats_acc: list[P3OSufficientStats] = [ + P3OSufficientStats.zeros(device=torch.cuda.current_device() if torch.cuda.is_available() else "cpu") + ] + invalid_count_acc = [stats_acc[0].valid_token_count.clone()] + + def forward_step( + iterator: DataIterator, + model_chunk: torch.nn.Module, + return_schedule_plan: bool = False, + ) -> tuple[torch.Tensor, callable]: + if return_schedule_plan: + raise ValueError("P3O ESS pre-pass does not support schedule plan generation") + batch = get_batch( + iterator, + [ + "tokens", + "multimodal_train_inputs", + "packed_seq_params", + "total_lengths", + "response_lengths", + "loss_masks", + "rollout_log_probs", + "max_seq_lens", + ], + args.data_pad_size_multiplier, + args.qkv_format, + args.allgather_cp, + getattr(args, "is_vl_model", False), + ) + # Use the same pure input builder as the training pass. The frozen cap + # must be calculated from exactly the logits used by the gradient pass. + forward_kwargs, needs_unsplit = build_rl_forward_kwargs(args, batch) + + # Dynamic CP: the VL bridge reads pg_collection.cp directly, so point it + # at this micro-batch's sub-group for the forward and restore after. + orig_cp_group = None + inner = None + dynamic_cp_size = batch.get("dynamic_cp_size") + if dynamic_cp_size is not None and needs_unsplit: + inner = model_chunk + while hasattr(inner, "module"): + inner = inner.module + orig_cp_group = inner.pg_collection.cp + inner.pg_collection.cp = mpu.get_dynamic_data_context_parallel_groups(group_size=dynamic_cp_size) + + try: + output_tensor = model_chunk(**forward_kwargs) + finally: + if orig_cp_group is not None: + inner.pg_collection.cp = orig_cp_group + + def collect(logits: torch.Tensor) -> tuple[torch.Tensor, int, dict[str, list | torch.Tensor]]: + # Only the pipeline last stage sees real logits; earlier stages just + # participate in the schedule. + if mpu.is_pipeline_last_stage(): + _, computed = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=batch["total_lengths"], + response_lengths=batch["response_lengths"], + with_entropy=False, + max_seq_lens=batch.get("max_seq_lens", None), + padded_total_lengths=batch.get("padded_total_lengths", None), + dynamic_cp_size=batch.get("dynamic_cp_size", None), + dynamic_cp_rank=batch.get("dynamic_cp_rank", None), + ) + # _local_stats_from_batch returns a device-resident invalid_flag + # instead of raising, so the non-finite detection rides the + # existing allreduce rather than adding a per-micro-batch + # GPU-CPU sync via bool() or .item(). + micro_stats, invalid_flag = _local_stats_from_batch(args, batch, computed["log_probs"]) + invalid_count_acc[0] = invalid_count_acc[0] + invalid_flag + stats_acc[0] = stats_acc[0] + micro_stats + zero = torch.zeros((), device=logits.device, dtype=torch.float32) + return zero, 1, {"keys": [], "values": zero.reshape(1)} + + return output_tensor, collect + + forward_backward_func = get_forward_backward_func() + + with preserved_iterator_positions(data_iterator), preserved_rng_state(), torch.no_grad(): + forward_backward_func( + forward_step_func=forward_step, + data_iterator=data_iterator, + model=model, + num_microbatches=num_microbatches, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + decoder_seq_length=args.decoder_seq_length, + forward_only=True, + ) + + # Accumulate every local micro-batch first, reduce exactly once over DP x CP + # on pipeline-last, then broadcast that fixed vector over PP. + distributed = torch.distributed.is_available() and torch.distributed.is_initialized() + is_pipeline_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=True) + dp_cp_group = ( + mpu.get_data_parallel_group(with_context_parallel=True) if distributed and is_pipeline_last_stage else None + ) + pp_group = ( + mpu.get_pipeline_model_parallel_group() + if distributed and mpu.get_pipeline_model_parallel_world_size() > 1 + else None + ) + reduced = synchronize_p3o_stats( + stats_acc[0], + invalid_count_acc[0], + dp_cp_group=dp_cp_group, + pp_group=pp_group, + is_pipeline_last_stage=is_pipeline_last_stage, + ) + step_context = finalize_p3o_step_context(reduced) + + if step_context.clamp_events: + logger.warning("P3O: clamped %d out-of-range ESS value(s) this step", step_context.clamp_events) + + return step_context + + +@contextmanager +def p3o_step_context_published(args: Namespace, step_context: P3OStepContext) -> Iterator[None]: + """Publish the step context on ``args`` for the duration of the train pass. + + The loss function reads the cap from here rather than from the micro-batch + dict: a per-micro-batch copy could diverge, and the whole point is that all + micro-batches of the step share one immutable cap. Cleared afterwards so a + stale cap can never leak into the next step. + """ + previous = getattr(args, P3O_STEP_CONTEXT_ATTR, None) + setattr(args, P3O_STEP_CONTEXT_ATTR, step_context) + try: + yield + finally: + setattr(args, P3O_STEP_CONTEXT_ATTR, previous) diff --git a/relax/backends/megatron/rollout_policy_lag.py b/relax/backends/megatron/rollout_policy_lag.py new file mode 100644 index 000000000..2c67030e2 --- /dev/null +++ b/relax/backends/megatron/rollout_policy_lag.py @@ -0,0 +1,102 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Scheduling helpers for periodic rollout policy snapshots.""" + +from typing import Protocol + + +ROLLOUT_POLICY_TAG = "rollout_policy" + + +class _TensorBackuperLike(Protocol): + def copy(self, *, src_tag: str, dst_tag: str) -> None: + """Copy one stored tensor snapshot to another tag.""" + + +def validate_update_weights_interval(update_weights_interval: int) -> int: + """Validate and return the rollout weight-update interval.""" + if update_weights_interval < 1: + raise ValueError(f"update_weights_interval must be a positive integer, got {update_weights_interval}") + return update_weights_interval + + +def rollout_weights_tag(update_weights_interval: int) -> str: + """Return the TensorBackuper tag whose weights should be pushed to + rollout.""" + interval = validate_update_weights_interval(update_weights_interval) + return ROLLOUT_POLICY_TAG if interval > 1 else "actor" + + +def compute_rollout_policy_age_rollouts( + current_rollout_id: int, + snapshot_rollout_id: int, +) -> int: + """Return the age of the behavior snapshot used by a training batch. + + The metric is emitted before the post-batch rollout snapshot refresh. At a + refresh boundary the just-trained batch therefore still reports the age of + the snapshot that generated it; the next batch observes the refreshed + snapshot. + """ + if current_rollout_id < 0: + raise ValueError("current_rollout_id must be non-negative") + if snapshot_rollout_id < 0: + raise ValueError("snapshot_rollout_id must be non-negative") + if current_rollout_id < snapshot_rollout_id: + raise ValueError("current_rollout_id cannot precede snapshot_rollout_id") + return current_rollout_id - snapshot_rollout_id + + +def initial_rollout_policy_snapshot_rollout(start_rollout_id: int) -> int: + """Return the snapshot version aligned with a fresh or resumed run.""" + if start_rollout_id < 0: + raise ValueError("start_rollout_id must be non-negative") + return start_rollout_id + + +def build_rollout_policy_age_metrics( + *, + current_rollout_id: int, + rollout_policy_snapshot_rollout: int, +) -> dict[str, int]: + """Build rollout-unit policy-age metrics for one training batch.""" + return { + "train/current_rollout_id": current_rollout_id, + "train/rollout_policy_snapshot_rollout": rollout_policy_snapshot_rollout, + "train/p3o/rollout_policy_age_rollouts": compute_rollout_policy_age_rollouts( + current_rollout_id, + rollout_policy_snapshot_rollout, + ), + } + + +def should_refresh_rollout_policy( + rollout_id: int, + update_weights_interval: int, + num_rollout: int, +) -> bool: + """Return whether the fixed rollout snapshot should adopt the trained + actor. + + The final step always refreshes so end-of-training evaluation sees the + latest actor even when the step is not an interval boundary. + """ + interval = validate_update_weights_interval(update_weights_interval) + completed_steps = rollout_id + 1 + return interval == 1 or completed_steps % interval == 0 or completed_steps == num_rollout + + +def maybe_refresh_rollout_policy( + weights_backuper: _TensorBackuperLike, + rollout_id: int, + update_weights_interval: int, + num_rollout: int, +) -> bool: + """Refresh a fixed rollout snapshot when its schedule reaches a + boundary.""" + interval = validate_update_weights_interval(update_weights_interval) + if interval == 1 or not should_refresh_rollout_policy(rollout_id, interval, num_rollout): + return False + + weights_backuper.copy(src_tag="actor", dst_tag=ROLLOUT_POLICY_TAG) + return True diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 751db235a..9bd13fc97 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -18,6 +18,7 @@ consume_opd_advantage_data, ) from relax.utils.training.ppo_utils import ( + GRPO_STYLE_ADVANTAGE_ESTIMATORS, compute_approx_kl, get_advantages_and_returns_batch, get_grpo_returns, @@ -172,7 +173,9 @@ def compute_advantages_and_returns(self, rollout_data: Dict[str, Any]) -> Dict[s for i in range(len(log_probs)) ] - if self.config.advantage_estimator in ["grpo", "gspo", "sapo", "cispo"]: + if self.config.advantage_estimator in GRPO_STYLE_ADVANTAGE_ESTIMATORS: + # P3O shares GRPO's group-relative advantage; the two differ only in + # how the policy-gradient coefficient is formed at loss time. rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device) returns = get_grpo_returns(rewards, kl) advantages = list(returns) # make a copy diff --git a/relax/core/registry.py b/relax/core/registry.py index 449609397..d2eb5d79d 100644 --- a/relax/core/registry.py +++ b/relax/core/registry.py @@ -88,6 +88,13 @@ class ROLES_PPO_FULLY_ASYNC_ON_POLICY(StrEnum): ROLES.reference: ActorFwd, ROLES.actor_fwd: ActorFwd, }, + "p3o": { + ROLES.rollout: Rollout, + ROLES.actor: Actor, + ROLES.advantages: Advantages, + ROLES.reference: ActorFwd, + ROLES.actor_fwd: ActorFwd, + }, "gspo": { ROLES.rollout: Rollout, ROLES.actor: Actor, diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py index 0a0149cda..035148aad 100644 --- a/relax/engine/rollout/sglang_rollout.py +++ b/relax/engine/rollout/sglang_rollout.py @@ -28,6 +28,7 @@ from relax.utils.async_utils import run from relax.utils.data.data import Dataset from relax.utils.data.processing_utils import ( + _sanitize_response_tokens_for_logprobs, async_encode_audio_for_rollout_engine, async_encode_image_for_rollout_engine, async_encode_video_tensor_for_rollout_engine, @@ -414,44 +415,21 @@ async def generate( output["meta_info"], new_response_tokens, new_response_log_probs ) - while hasattr(state.tokenizer, "image_token_id") and state.tokenizer.image_token_id in new_response_tokens: - index = new_response_tokens.index(state.tokenizer.image_token_id) - new_response_tokens[index] = state.tokenizer.pad_token_id - logger.warning( - "Image token found in output tokens, replaced with pad_token_id. Consider updating the model's stop condition to stop at image_token_id if you want to avoid this." - ) - - while hasattr(state.tokenizer, "audio_token_id") and state.tokenizer.audio_token_id in new_response_tokens: - index = new_response_tokens.index(state.tokenizer.audio_token_id) - new_response_tokens[index] = state.tokenizer.pad_token_id - logger.warning( - "Audio token found in output tokens, replaced with pad_token_id. Consider updating the model's stop condition to stop at audio_token_id if you want to avoid this." + if len(new_response_log_probs) > 0 and len(new_response_log_probs) != len(new_response_tokens): + raise ValueError( + "rollout response token/log-prob length mismatch: " + f"{len(new_response_tokens)} tokens vs {len(new_response_log_probs)} log-probs" ) - while hasattr(state.tokenizer, "video_token_id") and state.tokenizer.video_token_id in new_response_tokens: - index = new_response_tokens.index(state.tokenizer.video_token_id) - new_response_tokens[index] = state.tokenizer.pad_token_id + new_response_tokens, new_rollout_log_probs_mask, replacement_counts = _sanitize_response_tokens_for_logprobs( + state.tokenizer, state.processor, new_response_tokens + ) + for label, replaced in replacement_counts.items(): logger.warning( - "Video token found in output tokens, replaced with pad_token_id. Consider updating the model's stop condition to stop at video_token_id if you want to avoid this." + f"Replaced {replaced} stray {label} token(s) in rollout response with pad_token_id; " + "the corresponding behavior log-probs will be masked from training." ) - # K2.x tokenizers don't expose image_token_id but reserve <|media_pad|> - # for vision input slots. A hallucinated <|media_pad|> in the response - # inflates num_placeholders past sum(feature_lengths) in the bridge, - # forcing dynamic expansion → broadcast → 233 GiB OOM. Replace in-place - # so positional accounting matches sglang's per-token logprobs. - if state.processor is not None: - from relax.utils.data.processing_utils import sanitize_kimi_k25_response_tokens - - sanitized = sanitize_kimi_k25_response_tokens(state.processor, new_response_tokens) - if sanitized is not new_response_tokens: - replaced = sum(1 for a, b in zip(new_response_tokens, sanitized, strict=True) if a != b) - if replaced: - logger.warning( - f"K2.x: replaced {replaced} stray <|media_pad|> token(s) in rollout response with pad_token_id." - ) - new_response_tokens = sanitized - # Update sample with tokens directly - avoiding re-tokenization sample.tokens = sample.tokens + new_response_tokens sample.rollout_tokens = sample.rollout_tokens + new_response_tokens @@ -463,9 +441,23 @@ async def generate( assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout sample.loss_mask += [1] * len(new_response_tokens) - if sample.rollout_log_probs is None: - sample.rollout_log_probs = [] - sample.rollout_log_probs += new_response_log_probs + if len(new_response_log_probs) > 0: + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + if sample.rollout_log_probs_mask is None: + sample.rollout_log_probs_mask = [True] * len(sample.rollout_log_probs) + sample.rollout_log_probs += new_response_log_probs + sample.rollout_log_probs_mask += new_rollout_log_probs_mask + else: + if sample.rollout_log_probs: + raise ValueError("rollout log-probs disappeared during a multi-turn response") + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + + if sample.rollout_log_probs_mask is not None and len(sample.rollout_log_probs_mask) != len( + sample.rollout_log_probs + ): + raise ValueError("accumulated rollout log-prob mask is not aligned with rollout log-probs") if state.opd_manager and not evaluation: state.opd_manager.after_rollout(sample, output) diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 3bf40c090..1972a1077 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -1735,6 +1735,7 @@ def add_algo_arguments(parser): "ppo", "sapo", "cispo", + "p3o", ], default="grpo", help=( @@ -1834,6 +1835,30 @@ def add_algo_arguments(parser): "If not set, we will use the logprobs from the actor model." ), ) + parser.add_argument( + "--p3o-ess-scope", + choices=["micro-batch", "step"], + default="micro-batch", + help="P3O ESS scope: paper-compatible micro-batch (default) or optimizer step.", + ) + parser.add_argument( + "--p3o-kl-mode", + choices=["proxy", "proxy_safe"], + default="proxy", + help="P3O behavior-KL implementation: proxy or proxy_safe.", + ) + parser.add_argument( + "--clip-low", + type=float, + default=0.2, + help="Lower ratio margin used only for P3O clip-fraction monitoring.", + ) + parser.add_argument( + "--clip-high", + type=float, + default=0.2, + help="Upper ratio margin used only for P3O clip-fraction monitoring.", + ) # Off-Policy Correction using Importance Sampling: https://fengyao.notion.site/off-policy-rl parser.add_argument( "--use-tis", @@ -2872,6 +2897,92 @@ def _validate_agentic_rollout_args(args) -> None: raise ValueError("--agentic-eval-prepare-pool-size must be > 0.") +def _validate_p3o_args(args: argparse.Namespace) -> None: + """Reject P3O configurations whose ESS scope or replay would be wrong. + + These are hard errors, not warnings. Every condition below silently changes + the objective (not just performance), and the failure mode is a plausible + loss curve that does not implement P3O. + """ + # These are raises rather than asserts on purpose: `python -O` strips + # asserts, and every condition here silently changes the objective rather + # than crashing, so a stripped check would let a non-P3O run masquerade as + # one for its entire duration. + scope = getattr(args, "p3o_ess_scope", "micro-batch") + if scope not in {"micro-batch", "step"}: + raise ValueError(f"--p3o-ess-scope must be micro-batch or step, got {scope!r}.") + kl_mode = getattr(args, "p3o_kl_mode", "proxy") + if kl_mode not in {"proxy", "proxy_safe"}: + raise ValueError(f"--p3o-kl-mode must be proxy or proxy_safe, got {kl_mode!r}.") + clip_low = getattr(args, "clip_low", 0.2) + clip_high = getattr(args, "clip_high", 0.2) + if clip_low < 0.0 or clip_high < 0.0: + raise ValueError(f"--clip-low/--clip-high must be non-negative, got {clip_low}, {clip_high}.") + + if not args.use_rollout_logprobs: + raise ValueError( + "P3O requires the rollout sampling distribution as its behavior policy. " + "Add --use-rollout-logprobs; without it there is no importance ratio to correct." + ) + if not args.calculate_per_token_loss: + raise ValueError( + "P3O requires --calculate-per-token-loss. Per-sample-mean normalization " + "reintroduces a per-micro-batch denominator, so the loss would depend on " + "how the optimizer step is split into micro-batches." + ) + if args.use_tis: + raise ValueError( + "P3O and TIS (--use-tis) are mutually exclusive: both correct the same " + "rollout/training mismatch, and stacking them double-corrects the ratio." + ) + if getattr(args, "use_critic", False): + raise ValueError( + "P3O does not use a critic; it is a score-function estimator over group-relative " + "advantages. Drop --use-critic." + ) + + incompatible_flags = { + "get_mismatch_metrics": "--get-mismatch-metrics", + "use_opsm": "--use-opsm", + "enable_mtp_training": "--enable-mtp-training", + "use_routing_replay": "--use-routing-replay", + "use_rollout_routing_replay": "--use-rollout-routing-replay", + "overlap_moe_expert_parallel_comm": "--overlap-moe-expert-parallel-comm", + } + for attr, flag in incompatible_flags.items(): + if getattr(args, attr, False): + raise ValueError(f"P3O does not support {flag} in the replayed two-pass optimizer step.") + if getattr(args, "custom_pg_loss_reducer_function_path", None) is not None: + raise ValueError("P3O requires token-sum normalization and does not support a custom PG-loss reducer.") + + if scope == "step": + # The ESS pre-pass replays the same micro-batch window under no_grad. Ops + # that mutate state on a forward would make the two passes disagree. + if getattr(args, "fp8", None) is not None: + raise ValueError( + "P3O's ESS pre-pass runs a second forward over the same window, which would " + "advance FP8 amax history and make the training forward non-reproducible. " + "Disable FP8 or use --p3o-ess-scope micro-batch." + ) + dropout = max( + getattr(args, "attention_dropout", 0.0) or 0.0, + getattr(args, "hidden_dropout", 0.0) or 0.0, + (getattr(args, "lora_dropout", 0.0) or 0.0) if getattr(args, "lora_rank", 0) > 0 else 0.0, + ) + if dropout > 0.0: + raise ValueError( + f"P3O step scope requires deterministic replay, but dropout is enabled (max rate {dropout}). " + "Set attention, hidden, and LoRA dropout rates to 0.0 or use --p3o-ess-scope micro-batch." + ) + + if getattr(args, "fully_async", False): + raise ValueError( + "P3O's optimizer-step ESS scope requires the whole micro-batch window to be " + "available before the training pass. Fully-async mode streams micro-batches; " + "use --p3o-ess-scope micro-batch instead." + ) + + def _validate_reinforce_plus_plus_args(args, is_sft: bool) -> None: """Validate the frozen Task 29 REINFORCE++ algorithm contracts.""" if is_sft: @@ -3582,3 +3693,10 @@ def slime_validate_args(args): if args.genrm_model_path: args.genrm_engine_config = args.genrm_engine_config or {} args.genrm_sampling_config = args.genrm_sampling_config or {} + + # Validate the final effective values. Several execution flags are derived + # above (hybrid and routing replay), and custom YAML is applied near the end; + # validating earlier would let those paths silently bypass P3O's replay + # contract. + if args.advantage_estimator == "p3o": + _validate_p3o_args(args) diff --git a/relax/utils/data/processing_utils.py b/relax/utils/data/processing_utils.py index 73e2992d8..8dc150b59 100644 --- a/relax/utils/data/processing_utils.py +++ b/relax/utils/data/processing_utils.py @@ -302,6 +302,51 @@ def sanitize_kimi_k25_response_tokens( return [replacement_id if t == placeholder_id else t for t in response_tokens] +def _sanitize_response_tokens_for_logprobs( + tokenizer: object, + processor: object | None, + response_tokens: list[int], +) -> tuple[list[int], list[bool], dict[str, int]]: + """Replace invalid multimodal output tokens and mark stale log-prob + pairs.""" + sanitized = list(response_tokens) + pairing_mask = [True] * len(sanitized) + replacement_counts: dict[str, int] = {} + pad_token_id = int(getattr(tokenizer, "pad_token_id", 0) or 0) + + for label, attribute in ( + ("image", "image_token_id"), + ("audio", "audio_token_id"), + ("video", "video_token_id"), + ): + special_token_id = getattr(tokenizer, attribute, None) + if special_token_id is None: + continue + replaced = 0 + for index, token_id in enumerate(sanitized): + if token_id == special_token_id: + sanitized[index] = pad_token_id + pairing_mask[index] = False + replaced += 1 + if replaced: + replacement_counts[label] = replaced + + if processor is not None: + media_sanitized = sanitize_kimi_k25_response_tokens(processor, sanitized) + if len(media_sanitized) != len(sanitized): + raise ValueError("multimodal response sanitization must preserve token count") + replaced = 0 + for index, (before, after) in enumerate(zip(sanitized, media_sanitized, strict=True)): + if before != after: + pairing_mask[index] = False + replaced += 1 + if replaced: + replacement_counts["media_pad"] = replaced + sanitized = media_sanitized + + return sanitized, pairing_mask, replacement_counts + + def expand_kimi_k25_placeholders( processor: object, prompt_ids: list[int], diff --git a/relax/utils/opd/opd_utils.py b/relax/utils/opd/opd_utils.py index ce052263a..8c67d5f90 100644 --- a/relax/utils/opd/opd_utils.py +++ b/relax/utils/opd/opd_utils.py @@ -680,10 +680,28 @@ def add_opd_arguments(parser: Any) -> Any: return parser +def validate_p3o_opd_compatibility(args: Namespace) -> None: + """Reject the unsupported hybrid of P3O and on-policy distillation. + + P3O owns its behavior-policy correction, adaptive cap, and trust-region + loss. OPD can independently modify rollout payloads, advantages, or add a + teacher loss, so composing the two would optimize an objective that neither + implementation defines. + """ + if getattr(args, "advantage_estimator", None) == "p3o" and getattr(args, "use_opd", False): + raise ValueError( + "P3O and OPD are mutually exclusive: --advantage-estimator p3o uses an independent " + "policy-loss dispatch, while --use-opd changes teacher data, advantages, or loss terms. " + "Disable --use-opd or select a non-P3O advantage estimator." + ) + + def validate_opd_args(args: Namespace, *, is_sft: bool, log: Any = logger) -> None: if is_sft: return + validate_p3o_opd_compatibility(args) + if not getattr(args, "use_opd", False): return diff --git a/relax/utils/training/data_fields.py b/relax/utils/training/data_fields.py index ebd6479af..0ec869fb7 100644 --- a/relax/utils/training/data_fields.py +++ b/relax/utils/training/data_fields.py @@ -15,6 +15,8 @@ def _base_rollout_fields(args: Namespace) -> list[str]: ] if getattr(args, "use_rollout_routing_replay", False): fields.append("rollout_routed_experts") + if getattr(args, "use_rollout_logprobs", False): + fields.append("rollout_log_probs_mask") if getattr(args, "multimodal_keys", None) is not None: fields.append("multimodal_train_inputs") return fields diff --git a/relax/utils/training/p3o_replay.py b/relax/utils/training/p3o_replay.py new file mode 100644 index 000000000..cfcdbb1d4 --- /dev/null +++ b/relax/utils/training/p3o_replay.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Replay guards for P3O's two-pass optimizer step. + +P3O computes ESS over a whole optimizer step, so the data window must be read +twice: once to accumulate the importance-ratio moments, once to train. These two +context managers make the second read identical to what a single-pass run would +have seen -- same tokens, same RNG stream. They are deliberately free of any +Megatron import so the invariants can be tested on CPU. +""" + +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from typing import Any + +import torch + + +@contextmanager +def preserved_rng_state() -> Iterator[None]: + """Snapshot and restore CPU / CUDA / Megatron RNG around the stats pass. + + The train pass must see exactly the RNG stream it would have seen without a + pre-pass, otherwise any stochastic op (dropout, MoE jitter) would + desynchronize the two forwards -- and under tensor parallelism, the ranks + within one forward. + """ + cpu_state = torch.get_rng_state() + cuda_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None + + tracker = None + tracker_states = None + try: + from megatron.core.tensor_parallel.random import get_cuda_rng_tracker + + tracker = get_cuda_rng_tracker() + tracker_states = tracker.get_states() + except (ImportError, AssertionError, RuntimeError): + # Tracker unavailable or uninitialized (CPU tests, no model-parallel init). + tracker = None + + try: + yield + finally: + torch.set_rng_state(cpu_state) + if cuda_state is not None: + torch.cuda.set_rng_state(cuda_state) + if tracker is not None and tracker_states is not None: + tracker.set_states(tracker_states) + + +@contextmanager +def preserved_iterator_positions(data_iterator: Sequence[Any] | Any) -> Iterator[None]: + """Snapshot and restore data-iterator offsets, deduplicated by identity. + + Under virtual pipeline parallelism the same iterator instance is passed once + per model chunk. Restoring it twice would be harmless, but snapshotting it + twice and restoring in the wrong order would not, so dedupe on ``id``. + + The restore runs in ``finally``: a pre-pass that raises must still leave the + window replayable, so the error surfaces as itself rather than as a confusing + downstream shape mismatch. + + Raises: + RuntimeError: If an iterator cannot report its position, which would + silently make the train pass consume different tokens. + """ + iterators = data_iterator if isinstance(data_iterator, (list, tuple)) else [data_iterator] + + unique: dict[int, Any] = {} + for iterator in iterators: + if iterator is not None: + unique.setdefault(id(iterator), iterator) + + for iterator in unique.values(): + if not (hasattr(iterator, "snapshot_position") and hasattr(iterator, "restore_position")): + raise RuntimeError( + f"P3O: data iterator {type(iterator).__name__} is not replayable (missing " + "snapshot_position/restore_position). The optimizer-step ESS pre-pass must read " + "the window twice; materialize the window or disable --advantage-estimator p3o." + ) + + positions = {key: iterator.snapshot_position() for key, iterator in unique.items()} + try: + yield + # WARNING: callers must not advance or otherwise mutate any of the + # tracked iterators *outside* this context manager while the with-block + # is open. External advancement between snapshot and restore will + # silently corrupt the replay: restore_position rewinds to the saved + # offset, causing the train pass to re-consume tokens that were already + # consumed by the external caller rather than the tokens this pre-pass + # saw. Only the pre-pass (the model forward) should drive the iterators + # while this context is live. + finally: + for key, iterator in unique.items(): + iterator.restore_position(positions[key]) diff --git a/relax/utils/training/p3o_utils.py b/relax/utils/training/p3o_utils.py new file mode 100644 index 000000000..261fe4cca --- /dev/null +++ b/relax/utils/training/p3o_utils.py @@ -0,0 +1,474 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Pure-PyTorch primitives for P3O (adaptive policy optimization). + +P3O replaces PPO/GRPO's fixed clip range with a one-sided cap derived from the +normalized Effective Sample Size (ESS) of the token-level importance ratios, +and adds an adaptive trust region weighted by ``(1 - ESS)``. + +Reference: Fakoor et al., "Trust the Batch, On- or Off-Policy: Adaptive Policy +Optimization for RL Post-Training" (arXiv:2605.12380), Eq. (7), (11), (12) and +Appendix Algorithm 2. + +This module is deliberately free of any Megatron / ``mpu`` dependency: it owns +the formulas, the masking discipline and the stop-gradient boundaries, while +collectives and lifecycle live in the Megatron backend. The same sufficient +statistics support both paper-compatible micro-batch ESS and Relax's optional +optimizer-step ESS. +""" + +from dataclasses import dataclass + +import torch + + +# Epsilon placed in the ESS denominator. Kept bit-compatible with the reference +# implementation (FeynRL ``algs/P3O/p3o.py::calculate_ess``) so golden-value +# parity holds; intentionally not exposed as a CLI hyper-parameter. +ESS_DENOM_EPS = 1e-8 + +# Clamp applied to the exponent of the behavior-KL proxy, matching the reference +# (FeynRL ``algs/RL/common.py::compute_kl_distance``). +BEHAVIOR_KL_EXP_CLAMP = 10.0 + +# Shared by the checked and unchecked sufficient-statistics paths so the message +# a user sees does not depend on which one detected the non-finite ratio. +NONFINITE_RATIO_MESSAGE = ( + "P3O: non-finite importance ratio at a valid response token; refusing to " + "silently fall back to ESS=1. Check rollout log-probs and mask alignment." +) + + +def _require_identical_shapes(**tensors: torch.Tensor) -> None: + """Reject broadcasting between token-aligned P3O inputs.""" + shapes = {name: tuple(tensor.shape) for name, tensor in tensors.items()} + if len(set(shapes.values())) != 1: + formatted = ", ".join(f"{name}={shape}" for name, shape in shapes.items()) + raise ValueError(f"P3O token tensors must have identical shapes; got {formatted}") + + +@dataclass(frozen=True) +class P3OSufficientStats: + """Local (this-rank, this-micro-batch) ESS sufficient statistics. + + All three fields are ``float64`` scalar tensors so they can be stacked and + summed by a single collective without precision loss. + + Attributes: + sum_ratio: ``S1 = sum(rho_i)`` over valid response tokens. + sum_ratio_sq: ``S2 = sum(rho_i ** 2)`` over valid response tokens. + valid_token_count: ``N``, the number of valid response tokens. + """ + + sum_ratio: torch.Tensor + sum_ratio_sq: torch.Tensor + valid_token_count: torch.Tensor + + def as_vector(self) -> torch.Tensor: + """Stack the statistics into a ``[3]`` float64 tensor for reduction.""" + return torch.stack([self.sum_ratio, self.sum_ratio_sq, self.valid_token_count]) + + @classmethod + def zeros(cls, device: torch.device | str = "cpu") -> "P3OSufficientStats": + """Return all-zero statistics, used for dummy micro-batches.""" + zero = torch.zeros((), dtype=torch.float64, device=device) + return cls(sum_ratio=zero.clone(), sum_ratio_sq=zero.clone(), valid_token_count=zero.clone()) + + @classmethod + def from_vector(cls, vector: torch.Tensor) -> "P3OSufficientStats": + """Rebuild statistics from a reduced ``[3]`` tensor.""" + if vector.numel() != 3: + raise ValueError(f"expected a 3-element stat vector, got shape {tuple(vector.shape)}") + flat = vector.reshape(3).to(torch.float64) + return cls(sum_ratio=flat[0], sum_ratio_sq=flat[1], valid_token_count=flat[2]) + + def __add__(self, other: "P3OSufficientStats") -> "P3OSufficientStats": + """Accumulate statistics across micro-batches on the same rank.""" + return P3OSufficientStats( + sum_ratio=self.sum_ratio + other.sum_ratio, + sum_ratio_sq=self.sum_ratio_sq + other.sum_ratio_sq, + valid_token_count=self.valid_token_count + other.valid_token_count, + ) + + +@dataclass(frozen=True) +class P3OStepContext: + """Immutable per-optimizer-step P3O state shared by every micro-batch. + + Attributes: + normalized_ess: Global normalized ESS in ``[0, 1]``. + adaptive_cap: The ratio cap. Numerically equal to ``normalized_ess`` but + kept separate because it plays a different role in the objective. + valid_token_count: Global valid response-token count ``N``. + ratio_mean: ``S1 / N``. + ratio_std: Population std derived from the global moments. + clamp_events: Compatibility field. ESS is clamped on-device without a + host synchronization, so this remains zero. + """ + + normalized_ess: torch.Tensor + adaptive_cap: torch.Tensor + valid_token_count: torch.Tensor + ratio_mean: torch.Tensor + ratio_std: torch.Tensor + clamp_events: int = 0 + + +@dataclass(frozen=True) +class P3OTokenTerms: + """Element-wise P3O loss terms for one micro-batch. + + Every tensor has the shape of the concatenated response tokens and carries + no reduction, so the caller applies its own masking / normalization. + + Attributes: + ratio: ``rho_i``, detached. + score_loss: ``-sg(min(rho_i, cap)) * log_prob_i * sg(A_i)``. + behavior_kl_proxy: k3-style sampled-token KL against the behavior + policy, *not* multiplied by ``(1 - ESS)``. Keeps gradient. + adaptive_kl_loss: ``(1 - ESS) * behavior_kl_proxy``. + cap_hits: 1.0 where ``rho_i > cap``, else 0.0. + clip_hits: 1.0 where ``rho_i`` is outside the monitoring interval. + """ + + ratio: torch.Tensor + score_loss: torch.Tensor + behavior_kl_proxy: torch.Tensor + adaptive_kl_loss: torch.Tensor + cap_hits: torch.Tensor + clip_hits: torch.Tensor + + +def compute_p3o_log_ratio( + log_probs: torch.Tensor, + behavior_log_probs: torch.Tensor, + valid_mask: torch.Tensor, +) -> torch.Tensor: + """Compute the masked log importance ratio ``l_i``. + + Invalid positions are zeroed *before* any exponentiation so that padded + entries holding ``inf`` / ``NaN`` cannot poison the statistics via + ``inf * 0 -> NaN``. + + Args: + log_probs: Current-policy log-probs of the sampled tokens. + behavior_log_probs: Log-probs under the policy that actually generated + the tokens (rollout log-probs), already detached by the caller. + valid_mask: Boolean mask selecting valid response tokens. + + Returns: + ``l_i = log pi_theta - log pi_b`` in float32, zero at invalid positions. + """ + _require_identical_shapes( + log_probs=log_probs, + behavior_log_probs=behavior_log_probs, + valid_mask=valid_mask, + ) + log_ratio = log_probs.float() - behavior_log_probs.float() + return torch.where(valid_mask, log_ratio, torch.zeros_like(log_ratio)) + + +def compute_p3o_sufficient_stats( + log_probs: torch.Tensor, + behavior_log_probs: torch.Tensor, + valid_mask: torch.Tensor, +) -> P3OSufficientStats: + """Accumulate this micro-batch's contribution to the global ESS. + + The statistics are computed in float64 and fully detached: ESS is a + stop-gradient quantity in the P3O objective. + + Args: + log_probs: Current-policy log-probs of the sampled tokens. + behavior_log_probs: Behavior-policy (rollout) log-probs. + valid_mask: Boolean mask selecting valid response tokens. Prompt, + padding, CP padding and masked tokens must already be excluded. + + Returns: + Local :class:`P3OSufficientStats` in float64. + + Raises: + ValueError: If a valid position produced a non-finite ratio. + """ + stats, invalid_flag = compute_p3o_sufficient_stats_unchecked(log_probs, behavior_log_probs, valid_mask) + # This convenience wrapper is used outside the micro-batch hot path, so an + # eager host check gives callers a deterministic error. Training uses the + # unchecked variant and reduces the device flag with the ESS moments. + if bool(invalid_flag > 0): + raise ValueError(NONFINITE_RATIO_MESSAGE) + return stats + + +def compute_p3o_sufficient_stats_unchecked( + log_probs: torch.Tensor, + behavior_log_probs: torch.Tensor, + valid_mask: torch.Tensor, +) -> tuple[P3OSufficientStats, torch.Tensor]: + """Sync-free variant: report non-finite ratios as a device-resident flag. + + Identical arithmetic to :func:`compute_p3o_sufficient_stats`, but the + finiteness verdict is returned as a ``float64`` scalar tensor instead of + being tested on the host. This is what the ESS pre-pass calls: it runs once + per micro-batch, and a ``bool()`` there would stall the GPU pipeline + ``num_microbatches`` times per optimizer step. The flag rides along with + ``S1/S2/N`` through the step's single all-reduce, so the error still + surfaces on every rank -- just one collective later. + + Args: + log_probs: Current-policy log-probs of the sampled tokens. + behavior_log_probs: Behavior-policy (rollout) log-probs. + valid_mask: Boolean mask selecting valid response tokens. + + Returns: + ``(stats, invalid_flag)``. ``invalid_flag`` is ``1.0`` when any valid + position produced a non-finite ratio, else ``0.0``. When it is set, the + statistics are zeroed so a caller that defers the check cannot poison + ``S1/S2`` with ``inf``/``nan`` in the meantime. + """ + with torch.no_grad(): + mask_bool = valid_mask.bool() + log_ratio = compute_p3o_log_ratio(log_probs.detach(), behavior_log_probs.detach(), mask_bool) + + ratio = torch.exp(log_ratio.to(torch.float64)) + ratio = torch.where(mask_bool, ratio, torch.zeros_like(ratio)) + + # Both checks stay on device. log_ratio is already zeroed outside the + # mask, so a global isfinite() over it is equivalent to masking first. + invalid_flag = (~(torch.isfinite(log_ratio).all() & torch.isfinite(ratio).all())).to(torch.float64) + + # Zero the contribution when invalid, so deferring the host-side check + # cannot let inf/nan reach the reduced moments. + keep = 1.0 - invalid_flag + return ( + P3OSufficientStats( + sum_ratio=ratio.sum() * keep, + sum_ratio_sq=ratio.pow(2).sum() * keep, + valid_token_count=mask_bool.sum().to(torch.float64) * keep, + ), + invalid_flag, + ) + + +def finalize_p3o_step_context(stats: P3OSufficientStats) -> P3OStepContext: + """Turn globally reduced sufficient statistics into a frozen step context. + + Implements the paper's ``e = sg(S1^2 / (N * S2))`` with the reference + implementation's epsilon placement, i.e. ``S1^2 / (N * (S2 + eps))``. + + Args: + stats: Sufficient statistics already summed across DP x CP. + + Returns: + Immutable :class:`P3OStepContext` reused by every micro-batch of the + current optimizer step. + + Non-finite statistics and an empty valid-token set use the reference + implementation's neutral fallback: ``ESS=cap=1``, ratio mean 1 and ratio + std 0. This stays device-resident and does not synchronize a CUDA hot path. + """ + sum_ratio = stats.sum_ratio.to(torch.float64) + sum_ratio_sq = stats.sum_ratio_sq.to(torch.float64) + count = stats.valid_token_count.to(torch.float64) + + valid = torch.stack((sum_ratio, sum_ratio_sq, count)).isfinite().all() & (count >= 0.5) + one = torch.ones((), dtype=torch.float64, device=count.device) + zero = torch.zeros((), dtype=torch.float64, device=count.device) + safe_sum_ratio = torch.where(valid, sum_ratio, one) + safe_sum_ratio_sq = torch.where(valid, sum_ratio_sq, one) + safe_count = torch.where(valid, count, one) + + raw_ess = safe_sum_ratio.pow(2) / (safe_count * (safe_sum_ratio_sq + ESS_DENOM_EPS)) + ess = torch.where(valid, raw_ess.clamp(min=0.0, max=1.0), one) + + ratio_mean = torch.where(valid, safe_sum_ratio / safe_count, one) + variance = (safe_sum_ratio_sq / safe_count) - ratio_mean.pow(2) + ratio_std = torch.where(valid, variance.clamp(min=0.0).sqrt(), zero) + valid_token_count = torch.where(torch.isfinite(count) & (count >= 0.0), count, zero) + + return P3OStepContext( + normalized_ess=ess, + adaptive_cap=ess.clone(), + valid_token_count=valid_token_count, + ratio_mean=ratio_mean, + ratio_std=ratio_std, + clamp_events=0, + ) + + +class _P3OProxySafeK3(torch.autograd.Function): + """FeynRL k3 forward with a bounded, sign-correct extreme backward.""" + + @staticmethod + def forward(ctx, log_ratio: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(log_ratio) + exponent = torch.clamp(-log_ratio, min=-BEHAVIOR_KL_EXP_CLAMP, max=BEHAVIOR_KL_EXP_CLAMP) + return log_ratio + torch.exp(exponent) - 1.0 + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]: + (log_ratio,) = ctx.saved_tensors + exponent = torch.clamp(-log_ratio, min=-BEHAVIOR_KL_EXP_CLAMP, max=BEHAVIOR_KL_EXP_CLAMP) + gradient = 1.0 - torch.exp(exponent) + return (grad_output * gradient,) + + +def compute_p3o_behavior_kl_proxy( + log_probs: torch.Tensor, + behavior_log_probs: torch.Tensor, + valid_mask: torch.Tensor, + mode: str = "proxy", +) -> torch.Tensor: + """Sampled-token k3 proxy for ``KL(pi_theta || pi_b)``. + + ``K_i = l_i + exp(clip(-l_i, -C, C)) - 1`` with ``l_i`` the log ratio and + ``C = BEHAVIOR_KL_EXP_CLAMP`` (currently 10). When ``|l_i| > C`` the + exponent saturates: for ``l_i > C`` the exp term floors at ``exp(-C)`` so + the gradient of the kl term w.r.t. ``log_probs`` approaches 1 (only the + ``l_i`` addend contributes); for ``l_i < -C`` it caps at ``exp(C)`` + preventing numerical overflow. + Gradient flows through ``log_probs``, which is what makes this an adaptive + trust region rather than a diagnostic. + + This is a *proxy*: replay only stores the sampled token's log-prob, so the + full-vocabulary KL of the paper is not recoverable here. Do not report it as + the exact paper quantity. + + Args: + log_probs: Current-policy log-probs of the sampled tokens. + behavior_log_probs: Behavior-policy (rollout) log-probs, detached. + valid_mask: Boolean mask selecting valid response tokens. + + mode: ``proxy`` preserves the FeynRL autograd behavior. ``proxy_safe`` + preserves the exact forward values but corrects the saturated + negative-log-ratio gradient direction. + + Returns: + Element-wise KL proxy, zero at invalid positions. + """ + if mode not in {"proxy", "proxy_safe"}: + raise ValueError(f"P3O sampled-token KL mode must be proxy or proxy_safe, got {mode!r}") + mask_bool = valid_mask.bool() + log_ratio = compute_p3o_log_ratio(log_probs, behavior_log_probs, mask_bool) + if mode == "proxy_safe": + kl = _P3OProxySafeK3.apply(log_ratio) + else: + exponent = torch.clamp(-log_ratio, min=-BEHAVIOR_KL_EXP_CLAMP, max=BEHAVIOR_KL_EXP_CLAMP) + kl = log_ratio + torch.exp(exponent) - 1.0 + return torch.where(mask_bool, kl, torch.zeros_like(kl)) + + +def compute_p3o_exact_kl( + policy_logits: torch.Tensor, + behavior_logits: torch.Tensor, + valid_mask: torch.Tensor, +) -> torch.Tensor: + """Compute the exact forward KL over a full vocabulary. + + This pure helper is intended for small-vocabulary verification and for a + future training path that carries behavior logits. Production rollout data + currently stores only selected-token log-probs, so the loss integration + rejects ``exact`` mode with an explicit error. + """ + if policy_logits.shape != behavior_logits.shape: + raise ValueError( + "P3O exact-KL logits must have identical shapes; " + f"got policy={tuple(policy_logits.shape)}, behavior={tuple(behavior_logits.shape)}" + ) + if policy_logits.ndim < 1 or tuple(policy_logits.shape[:-1]) != tuple(valid_mask.shape): + raise ValueError( + "P3O exact-KL valid_mask must match the logits token dimensions; " + f"got logits={tuple(policy_logits.shape)}, mask={tuple(valid_mask.shape)}" + ) + + mask_bool = valid_mask.bool() + expanded_mask = mask_bool.unsqueeze(-1) + safe_policy_logits = torch.where(expanded_mask, policy_logits.float(), torch.zeros_like(policy_logits.float())) + safe_behavior_logits = torch.where( + expanded_mask, + behavior_logits.detach().float(), + torch.zeros_like(behavior_logits.detach().float()), + ) + policy_log_probs = torch.log_softmax(safe_policy_logits, dim=-1) + behavior_log_probs = torch.log_softmax(safe_behavior_logits, dim=-1) + exact_kl = (policy_log_probs.exp() * (policy_log_probs - behavior_log_probs)).sum(dim=-1) + return torch.where(mask_bool, exact_kl, torch.zeros_like(exact_kl)) + + +def compute_p3o_token_terms( + log_probs: torch.Tensor, + behavior_log_probs: torch.Tensor, + advantages: torch.Tensor, + valid_mask: torch.Tensor, + step_context: P3OStepContext, + kl_mode: str = "proxy", + clip_low: float = 0.2, + clip_high: float = 0.2, +) -> P3OTokenTerms: + """Compute the element-wise P3O loss terms for one micro-batch. + + The score-function term is ``-sg(min(rho_i, cap)) * log pi_theta * sg(A_i)``. + The *entire* ``min(rho, cap)`` factor is detached, not just the cap: P3O is a + REINFORCE-style update whose only gradient path is ``log_probs``. There is no + lower cap and no advantage-sign-dependent branch, so ``eps_clip`` plays no + part in the objective. + + Args: + log_probs: Current-policy log-probs of the sampled tokens (gradient + source). + behavior_log_probs: Behavior-policy (rollout) log-probs. + advantages: GRPO group-relative advantages broadcast to response tokens. + valid_mask: Boolean mask selecting valid response tokens. + step_context: Frozen context carrying this optimizer step's global cap. + kl_mode: Sampled-token behavioral KL implementation. ``exact`` is + rejected because this function receives no behavior logits. + clip_low: Lower monitoring margin around ratio 1. + clip_high: Upper monitoring margin around ratio 1. + + Returns: + :class:`P3OTokenTerms` with no reduction applied. + """ + if kl_mode == "exact": + raise ValueError( + "P3O exact KL requires full-vocabulary behavior logits; rollout data currently stores only " + "selected-token log-probs. Use proxy/proxy_safe for training." + ) + if clip_low < 0.0 or clip_high < 0.0: + raise ValueError(f"P3O clip monitoring margins must be non-negative, got {clip_low}, {clip_high}") + + _require_identical_shapes( + log_probs=log_probs, + behavior_log_probs=behavior_log_probs, + advantages=advantages, + valid_mask=valid_mask, + ) + mask_bool = valid_mask.bool() + behavior_log_probs = behavior_log_probs.detach() + cap = step_context.adaptive_cap.to(dtype=torch.float32, device=log_probs.device) + ess = step_context.normalized_ess.to(dtype=torch.float32, device=log_probs.device) + + with torch.no_grad(): + log_ratio_detached = compute_p3o_log_ratio(log_probs.detach(), behavior_log_probs, mask_bool) + ratio = torch.exp(log_ratio_detached) + ratio = torch.where(mask_bool, ratio, torch.zeros_like(ratio)) + # Full stop-gradient on min(ratio, cap): the coefficient must not + # contribute a gradient path of its own. + # Keep the cap on device. Converting it with ``float(cap)`` would add a + # GPU-to-CPU synchronization in every training micro-batch. + coefficient = torch.minimum(ratio, cap) + cap_hits = (mask_bool & (ratio > cap)).to(dtype=torch.float32) + clip_hits = (mask_bool & ((ratio < 1.0 - clip_low) | (ratio > 1.0 + clip_high))).to(dtype=torch.float32) + + score_loss = -(coefficient * log_probs.float() * advantages.detach().float()) + score_loss = torch.where(mask_bool, score_loss, torch.zeros_like(score_loss)) + + behavior_kl_proxy = compute_p3o_behavior_kl_proxy(log_probs, behavior_log_probs, mask_bool, mode=kl_mode) + adaptive_kl_loss = (1.0 - ess) * behavior_kl_proxy + + return P3OTokenTerms( + ratio=ratio, + score_loss=score_loss, + behavior_kl_proxy=behavior_kl_proxy, + adaptive_kl_loss=adaptive_kl_loss, + cap_hits=cap_hits, + clip_hits=clip_hits, + ) diff --git a/relax/utils/training/ppo_utils.py b/relax/utils/training/ppo_utils.py index 4d7ba873d..768fb3fc2 100644 --- a/relax/utils/training/ppo_utils.py +++ b/relax/utils/training/ppo_utils.py @@ -17,6 +17,10 @@ logger = get_logger(__name__) +GRPO_STYLE_ADVANTAGE_ESTIMATORS = frozenset({"grpo", "gspo", "sapo", "cispo", "p3o"}) +GROUP_REWARD_NORMALIZATION_ESTIMATORS = GRPO_STYLE_ADVANTAGE_ESTIMATORS | {"reinforce_plus_plus_baseline"} + + def validate_ppo_config(config: Namespace) -> None: if getattr(config, "advantage_estimator", None) != "ppo": return diff --git a/relax/utils/training/train_dump_utils.py b/relax/utils/training/train_dump_utils.py index d556377ea..cf079a41f 100644 --- a/relax/utils/training/train_dump_utils.py +++ b/relax/utils/training/train_dump_utils.py @@ -196,6 +196,7 @@ def _sample_to_summary_record(sample, rollout_id: int, idx: int, dataset_name: s total_length = len(sample.tokens) if sample.tokens else 0 response_length = sample.response_length prompt_length = max(total_length - response_length, 0) + response_token_ids = list(sample.tokens[prompt_length:]) if sample.tokens else [] multimodal_stats = get_sample_multimodal_stats(sample) metadata = sample.metadata or {} record = { @@ -209,6 +210,7 @@ def _sample_to_summary_record(sample, rollout_id: int, idx: int, dataset_name: s "total_length": total_length, "prompt_token_count": prompt_length, "response_token_count": response_length, + "response_token_ids": response_token_ids, "total_token_count": total_length, "image_count": multimodal_stats["image_count"], "image_token_count": multimodal_stats["image_token_count"], @@ -217,6 +219,10 @@ def _sample_to_summary_record(sample, rollout_id: int, idx: int, dataset_name: s "status": sample.status.value if hasattr(sample.status, "value") else str(sample.status), "group_index": sample.group_index, } + if sample.rollout_log_probs is not None: + record["response_rollout_log_probs"] = list(sample.rollout_log_probs) + if sample.rollout_log_probs_mask is not None: + record["response_rollout_log_probs_mask"] = list(sample.rollout_log_probs_mask) if sample.label is not None: record["label"] = sample.label if sample.multimodal_inputs is not None: diff --git a/relax/utils/types.py b/relax/utils/types.py index 9c7cabb5b..4219cf70b 100644 --- a/relax/utils/types.py +++ b/relax/utils/types.py @@ -27,6 +27,7 @@ class Sample: loss_mask: list[int] | None = None weight_versions: list[str] = field(default_factory=list) rollout_log_probs: list[float] | None = None # Log probabilities from rollout engine + rollout_log_probs_mask: list[bool] | None = None # True where token and behavior log-prob still correspond rollout_routed_experts: list[list[int]] | None = None # Routed experts from rollout engine remove_sample: bool = False abort_count: int = 0 # Number of times this sample has been aborted diff --git a/relax/utils/utils.py b/relax/utils/utils.py index 92bca9407..ae5558c28 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -15,6 +15,10 @@ from relax.utils.env import Envs, validate_env from relax.utils.logging_utils import get_logger from relax.utils.misc import load_function +from relax.utils.training.ppo_utils import ( + GROUP_REWARD_NORMALIZATION_ESTIMATORS, + GRPO_STYLE_ADVANTAGE_ESTIMATORS, +) from relax.utils.types import Sample @@ -110,10 +114,43 @@ def convert_samples_to_train_data(args: Any, samples: list[Sample] | list[list[S "sample_indices": [sample.index for sample in samples], } + has_rollout_log_probs = [sample.rollout_log_probs is not None for sample in samples] + if any(has_rollout_log_probs) and not all(has_rollout_log_probs): + raise ValueError("rollout_log_probs must be present for every sample in a training batch or for none of them") + if getattr(args, "use_rollout_logprobs", False) and not all(has_rollout_log_probs): + raise ValueError("--use-rollout-logprobs requires behavior log-probs for every training sample") + + rollout_log_probs_masks: list[list[bool]] | None = None + if all(has_rollout_log_probs): + candidate_masks: list[list[bool]] = [] + masks_aligned = True + for sample in samples: + rollout_log_probs = sample.rollout_log_probs + assert rollout_log_probs is not None + if len(rollout_log_probs) != sample.response_length: + if getattr(args, "use_rollout_logprobs", False) or rollout_log_probs: + raise ValueError( + f"rollout log-prob length {len(rollout_log_probs)} != response length {sample.response_length}" + ) + masks_aligned = False + continue + pairing_mask = sample.rollout_log_probs_mask + if pairing_mask is None: + pairing_mask = [True] * sample.response_length + if len(pairing_mask) != len(rollout_log_probs): + raise ValueError( + "rollout log-prob mask length " + f"{len(pairing_mask)} != rollout log-prob length {len(rollout_log_probs)}" + ) + sample.rollout_log_probs_mask = [bool(value) for value in pairing_mask] + candidate_masks.append(sample.rollout_log_probs_mask) + if masks_aligned and len(candidate_masks) == len(samples): + rollout_log_probs_masks = candidate_masks + # loss mask # TODO: compress the loss mask loss_masks = [] - for sample in samples: + for sample_index, sample in enumerate(samples): # always instantiate loss_mask if not provided if sample.loss_mask is None: sample.loss_mask = [1] * sample.response_length @@ -126,6 +163,15 @@ def convert_samples_to_train_data(args: Any, samples: list[Sample] | list[list[S ) if sample.remove_sample: sample.loss_mask = [0] * sample.response_length + if rollout_log_probs_masks is not None: + sample.loss_mask = [ + int(bool(loss_value) and pairing_value) + for loss_value, pairing_value in zip( + sample.loss_mask, + rollout_log_probs_masks[sample_index], + strict=True, + ) + ] loss_masks.append(sample.loss_mask) train_data["loss_masks"] = loss_masks @@ -142,8 +188,10 @@ def convert_samples_to_train_data(args: Any, samples: list[Sample] | list[list[S train_data["round_number"] = [sample.metadata["round_number"] for sample in samples] # Add rollout log probabilities for off-policy correction - if samples[0].rollout_log_probs is not None: + if all(has_rollout_log_probs): train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples] + if rollout_log_probs_masks is not None: + train_data["rollout_log_probs_mask"] = rollout_log_probs_masks if samples[0].rollout_routed_experts is not None: train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples] @@ -181,10 +229,7 @@ def post_process_rewards(args: Any, samples: list[Sample] | list[list[Sample]]): raw_rewards = [sample.get_reward_value(args) for sample in samples] if getattr(args, "agentic_custom_advantage_path", None) is not None: return raw_rewards, [sample.custom_advantage for sample in samples] - if ( - args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline"] - and args.rewards_normalization - ): + if args.advantage_estimator in GROUP_REWARD_NORMALIZATION_ESTIMATORS and args.rewards_normalization: # group norm rewards = torch.tensor(raw_rewards, dtype=torch.float) positions_by_group: dict[int, list[int]] = {} @@ -202,9 +247,14 @@ def post_process_rewards(args: Any, samples: list[Sample] | list[list[Sample]]): f"Reward group {group_index} has {len(positions)} samples, expected {args.n_samples_per_prompt}." ) group_rewards = rewards[positions] - group_rewards = group_rewards - group_rewards.mean() - if args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo"] and args.grpo_std_normalization: - group_rewards = group_rewards / (group_rewards.std() + 1e-6) + if args.advantage_estimator == "p3o": + if len(positions) > 1: + group_rewards = group_rewards - group_rewards.mean() + group_rewards = group_rewards / (group_rewards.std(correction=1) + 1e-8) + else: + group_rewards = group_rewards - group_rewards.mean() + if args.advantage_estimator in GRPO_STYLE_ADVANTAGE_ESTIMATORS and args.grpo_std_normalization: + group_rewards = group_rewards / (group_rewards.std() + 1e-6) normalized_rewards[positions] = group_rewards return raw_rewards, normalized_rewards.tolist() @@ -437,7 +487,7 @@ def get_debug_data(args, rollout_id: int, batch_size, dp_rank: int) -> Dict[str, original_num_rows = len(data) if ( args.custom_reward_post_process_path is None - and args.advantage_estimator in ["grpo", "gspo", "sapo", "cispo", "reinforce_plus_plus_baseline"] + and args.advantage_estimator in GROUP_REWARD_NORMALIZATION_ESTIMATORS and args.rewards_normalization ): group_ids = list(dict.fromkeys(sample.group_index for sample in data)) diff --git a/scripts/models/qwen3-4B.sh b/scripts/models/qwen3-4B.sh index 747d4c652..baf110bba 100644 --- a/scripts/models/qwen3-4B.sh +++ b/scripts/models/qwen3-4B.sh @@ -12,7 +12,7 @@ MODEL_ARGS=( --disable-bias-linear --normalization "RMSNorm" --norm-epsilon 1e-6 - --rotary-base 1000000 + --rotary-base "${MODEL_ARGS_ROTARY_BASE:-1000000}" --vocab-size 151936 --kv-channels 128 --qk-layernorm diff --git a/tests/backends/megatron/_megatron_stub.py b/tests/backends/megatron/_megatron_stub.py new file mode 100644 index 000000000..350f04ee8 --- /dev/null +++ b/tests/backends/megatron/_megatron_stub.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Import-time megatron stubs for CPU-only P3O tests. + +``relax.backends.megatron.{loss,model,p3o_step}`` import ``megatron.core`` at +module scope, but CI installs no megatron package (see +``.github/workflows/ci.yml``). The P3O logic under test is pure tensor math plus +collectives, so the megatron surface can be replaced by ``MagicMock`` for the +duration of the import. + +Without this, the four P3O test modules raise ``ModuleNotFoundError`` during +collection, and because CI runs ``pytest tests/ -x`` that aborts the *entire* +suite rather than skipping a few tests. + +Stubbing only spans the ``with`` block: the previous ``sys.modules`` entries are +restored afterwards, so a real megatron install is never shadowed and these +tests exercise the same code path on a GPU machine. +""" + +from __future__ import annotations + +import sys +from collections.abc import Iterator +from contextlib import contextmanager +from importlib.abc import Loader, MetaPathFinder +from importlib.machinery import ModuleSpec +from types import ModuleType +from unittest.mock import MagicMock + + +#: Top-level packages replaced while the context manager is active. Any +#: submodule below these is synthesized on demand, so the P3O import chain does +#: not have to be enumerated here. +STUBBED_ROOTS = ("megatron",) + + +class _MagicModule(ModuleType): + """Module whose unknown attributes resolve to ``MagicMock``. + + A plain ``MagicMock`` cannot stand in for a package -- ``import a.b`` fails + with "is not a package" -- so a real module object is used and attribute + lookup is delegated to a mock. + """ + + def __init__(self, name: str) -> None: + super().__init__(name) + self.__path__: list[str] = [] + self._mock = MagicMock(name=name) + + def __getattr__(self, item: str) -> object: + if item.startswith("__") and item.endswith("__"): + raise AttributeError(item) + return getattr(self._mock, item) + + +class _StubLoader(Loader): + def create_module(self, spec: ModuleSpec) -> ModuleType: + return _MagicModule(spec.name) + + def exec_module(self, module: ModuleType) -> None: # noqa: D102 - nothing to execute + return None + + +class _StubFinder(MetaPathFinder): + """Resolve any ```` or ``.*`` name to a synthetic module.""" + + def __init__(self, roots: tuple[str, ...]) -> None: + self._roots = roots + + def find_spec(self, fullname: str, path: object = None, target: object = None) -> ModuleSpec | None: + root = fullname.split(".", 1)[0] + if root not in self._roots: + return None + return ModuleSpec(fullname, _StubLoader(), is_package=True) + + +@contextmanager +def stubbed_megatron_modules(roots: tuple[str, ...] = STUBBED_ROOTS) -> Iterator[None]: + """Make ``megatron`` importable as a stub, restoring prior state on exit. + + No-op for roots that are genuinely installed, so a GPU machine with real + megatron exercises the production import path unchanged. + """ + missing = tuple(root for root in roots if _is_missing(root)) + if not missing: + yield + return + + finder = _StubFinder(missing) + sys.meta_path.insert(0, finder) + created_before = set(sys.modules) + try: + yield + finally: + if finder in sys.meta_path: + sys.meta_path.remove(finder) + for name in set(sys.modules) - created_before: + if isinstance(sys.modules.get(name), _MagicModule): + del sys.modules[name] + + +def _is_missing(root: str) -> bool: + if root in sys.modules: + return False + try: + from importlib.util import find_spec + + return find_spec(root) is None + except (ImportError, ValueError, ModuleNotFoundError): + return True diff --git a/tests/backends/megatron/test_data_vpp.py b/tests/backends/megatron/test_data_vpp.py index c2f8980aa..3e37941f2 100644 --- a/tests/backends/megatron/test_data_vpp.py +++ b/tests/backends/megatron/test_data_vpp.py @@ -47,6 +47,78 @@ def test_vpp_microbatch_rounding_uses_ceil_multiple(monkeypatch): assert rounded.tolist() == [4, 4, 4, 8] +def test_build_rl_forward_kwargs_uses_packed_text_inputs(monkeypatch): + data_module = _load_data_module(monkeypatch) + tokens = object() + packed_seq_params = object() + full_loss_masks = object() + + forward_kwargs, needs_unsplit = data_module.build_rl_forward_kwargs( + Namespace(is_vl_model=False, uses_unsplit_forward=False), + { + "tokens": tokens, + "packed_seq_params": packed_seq_params, + "full_loss_masks": full_loss_masks, + }, + ) + + assert not needs_unsplit + assert forward_kwargs["input_ids"] is tokens + assert forward_kwargs["packed_seq_params"] is packed_seq_params + assert forward_kwargs["attention_mask"] is None + assert forward_kwargs["loss_mask"] is full_loss_masks + + +def test_build_rl_forward_kwargs_uses_vl_thd_bridge_inputs(monkeypatch): + data_module = _load_data_module(monkeypatch) + unsplit_tokens = object() + attention_mask = object() + vlm_packed_seq_params = object() + pixel_values = object() + + forward_kwargs, needs_unsplit = data_module.build_rl_forward_kwargs( + Namespace(is_vl_model=True, uses_unsplit_forward=False), + { + "tokens": object(), + "unsplit_tokens": unsplit_tokens, + "packed_seq_params": object(), + "vlm_packed_seq_params": vlm_packed_seq_params, + "unsplit_attention_mask": attention_mask, + "full_loss_masks": object(), + "multimodal_train_inputs": {"pixel_values": pixel_values}, + }, + ) + + assert needs_unsplit + assert forward_kwargs["input_ids"] is unsplit_tokens + assert forward_kwargs["attention_mask"] is attention_mask + assert forward_kwargs["packed_seq_params"] is vlm_packed_seq_params + assert forward_kwargs["loss_mask"] is None + assert forward_kwargs["pixel_values"] is pixel_values + + +def test_build_rl_forward_kwargs_matches_custom_multimodal_gate(monkeypatch): + data_module = _load_data_module(monkeypatch) + unsplit_tokens = object() + pixel_values = object() + + forward_kwargs, needs_unsplit = data_module.build_rl_forward_kwargs( + Namespace(is_vl_model=False, uses_unsplit_forward=False), + { + "tokens": object(), + "unsplit_tokens": unsplit_tokens, + "packed_seq_params": object(), + "full_loss_masks": object(), + "multimodal_train_inputs": {"pixel_values": pixel_values}, + }, + ) + + assert needs_unsplit + assert forward_kwargs["input_ids"] is unsplit_tokens + assert forward_kwargs["packed_seq_params"] is None + assert "pixel_values" not in forward_kwargs + + def test_rollout_minibatch_plan_derives_from_global_batch(monkeypatch): data_module = _load_data_module(monkeypatch) args = Namespace( diff --git a/tests/backends/megatron/test_p3o_cp_metadata.py b/tests/backends/megatron/test_p3o_cp_metadata.py new file mode 100644 index 000000000..ae734f3ff --- /dev/null +++ b/tests/backends/megatron/test_p3o_cp_metadata.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Fail-fast tests for P3O context-parallel metadata alignment.""" + +from collections.abc import Callable + +import pytest +import torch + +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules + + +with stubbed_megatron_modules(): + from relax.backends.megatron.cp_utils import ( + get_cp_local_num_tokens, + get_cp_local_valid_mask, + get_sum_of_sample_mean, + ) + + +CP_METADATA_CONSUMERS: tuple[Callable[..., object], ...] = ( + get_sum_of_sample_mean, + get_cp_local_num_tokens, + get_cp_local_valid_mask, +) + + +@pytest.mark.parametrize("consumer", CP_METADATA_CONSUMERS) +@pytest.mark.parametrize( + "mismatched_field", + ["total_lengths", "response_lengths", "loss_masks", "max_seq_lens", "padded_total_lengths"], +) +def test_p3o_cp_metadata_length_mismatch_fails(consumer, mismatched_field): + metadata = { + "total_lengths": [3, 3], + "response_lengths": [2, 2], + "loss_masks": [torch.ones(2), torch.ones(2)], + "max_seq_lens": [3, 3], + "padded_total_lengths": [4, 4], + } + metadata[mismatched_field] = metadata[mismatched_field][:-1] + + with pytest.raises(ValueError, match=rf"CP metadata lengths must match;.*{mismatched_field}=1"): + consumer(**metadata) diff --git a/tests/backends/megatron/test_p3o_distributed.py b/tests/backends/megatron/test_p3o_distributed.py new file mode 100644 index 000000000..60df65d25 --- /dev/null +++ b/tests/backends/megatron/test_p3o_distributed.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Real Gloo checks for P3O stats and objective synchronization.""" + +from __future__ import annotations + +import math +import os +import socket + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules + + +with stubbed_megatron_modules(("megatron", "ray", "tensordict")): + from relax.backends.megatron import p3o_step + from relax.backends.megatron.p3o_step import synchronize_p3o_stats + +from relax.utils.training.p3o_utils import ( + P3OSufficientStats, + compute_p3o_sufficient_stats, + compute_p3o_token_terms, + finalize_p3o_step_context, +) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _init_gloo(rank: int, world_size: int, port: int) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + + +def _nonfinite_worker(rank: int, world_size: int, port: int) -> None: + _init_gloo(rank, world_size, port) + try: + p3o_step.mpu.is_pipeline_last_stage = lambda ignore_virtual=True: True + p3o_step.mpu.get_data_parallel_group = lambda with_context_parallel=True: dist.group.WORLD + p3o_step.mpu.get_pipeline_model_parallel_world_size = lambda: 1 + + stats = ( + P3OSufficientStats.zeros() + if rank == 0 + else P3OSufficientStats.from_vector(torch.tensor([1.0, 1.0, 1.0], dtype=torch.float64)) + ) + invalid_count = torch.tensor(float(rank == 0), dtype=torch.float64) + + try: + synchronize_p3o_stats( + stats, + invalid_count, + dp_cp_group=dist.group.WORLD, + pp_group=None, + is_pipeline_last_stage=True, + ) + except ValueError as error: + assert "non-finite importance ratio" in str(error) + else: + raise AssertionError("every rank must fail after the synchronized invalid flag") + + healthy = torch.ones((), dtype=torch.float64) + dist.all_reduce(healthy) + assert healthy.item() == world_size + finally: + dist.destroy_process_group() + + +def _pipeline_worker(rank: int, world_size: int, port: int) -> None: + _init_gloo(rank, world_size, port) + try: + dp_groups = [dist.new_group([dp_rank]) for dp_rank in range(world_size)] + p3o_step.mpu.is_pipeline_last_stage = lambda ignore_virtual=True: rank == world_size - 1 + p3o_step.mpu.get_data_parallel_group = lambda with_context_parallel=True: dp_groups[rank] + p3o_step.mpu.get_pipeline_model_parallel_world_size = lambda: world_size + p3o_step.mpu.get_pipeline_model_parallel_group = lambda: dist.group.WORLD + + expected = torch.tensor([7.5, 21.25, 4.0], dtype=torch.float64) + stats = P3OSufficientStats.from_vector(expected) if rank == world_size - 1 else P3OSufficientStats.zeros() + + synchronized = synchronize_p3o_stats( + stats, + torch.zeros((), dtype=torch.float64), + dp_cp_group=dp_groups[rank] if rank == world_size - 1 else None, + pp_group=dist.group.WORLD, + is_pipeline_last_stage=rank == world_size - 1, + ) + + torch.testing.assert_close(synchronized.as_vector(), expected, rtol=0.0, atol=0.0) + finally: + dist.destroy_process_group() + + +def _partition_worker(rank: int, world_size: int, port: int) -> None: + _init_gloo(rank, world_size, port) + try: + singleton_groups = [dist.new_group([group_rank]) for group_rank in range(world_size)] + dp2_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + active_group = [dist.group.WORLD] + p3o_step.mpu.is_pipeline_last_stage = lambda ignore_virtual=True: True + p3o_step.mpu.get_data_parallel_group = lambda with_context_parallel=True: active_group[0] + p3o_step.mpu.get_pipeline_model_parallel_world_size = lambda: 1 + + behavior = torch.full((11,), -2.0) + ratios = (1.0, 2.0, 0.5, 4.0, 0.8, 1.4, 0.25, 3.0, 1.1, 0.6, 2.5) + log_probs_value = behavior + torch.tensor([math.log(value) for value in ratios]) + advantages = torch.tensor([1.0, -1.0, 0.5, 2.0, -0.2, 0.7, -1.5, 1.2, 0.4, -0.8, 1.8]) + valid_mask = torch.tensor([True, True, False, True, True, False, True, True, True, False, True]) + all_indices = torch.arange(log_probs_value.numel()) + + oracle_context = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs_value, behavior, valid_mask)) + oracle_log_probs = log_probs_value.clone().requires_grad_(True) + oracle_terms = compute_p3o_token_terms( + oracle_log_probs, + behavior, + advantages, + valid_mask, + oracle_context, + ) + oracle_loss = (oracle_terms.score_loss + oracle_terms.adaptive_kl_loss).sum() + oracle_loss = oracle_loss / oracle_context.valid_token_count + oracle_loss.backward() + oracle_gradient = oracle_log_probs.grad.detach() + + def assert_partition(shards: list[torch.Tensor], process_group) -> None: + active_group[0] = process_group + local_stats = P3OSufficientStats.zeros() + for shard in shards: + if shard.numel() == 0: + local_stats = local_stats + P3OSufficientStats.zeros() + else: + local_stats = local_stats + compute_p3o_sufficient_stats( + log_probs_value[shard], + behavior[shard], + valid_mask[shard], + ) + synchronized = synchronize_p3o_stats( + local_stats, + torch.zeros((), dtype=torch.float64), + dp_cp_group=process_group, + pp_group=None, + is_pipeline_last_stage=True, + ) + context = finalize_p3o_step_context(synchronized) + torch.testing.assert_close(context.normalized_ess, oracle_context.normalized_ess) + + local_log_probs = log_probs_value.clone().requires_grad_(True) + local_total = 0.0 * local_log_probs.sum() + for shard in shards: + if shard.numel() == 0: + continue + terms = compute_p3o_token_terms( + local_log_probs[shard], + behavior[shard], + advantages[shard], + valid_mask[shard], + context, + ) + local_total = local_total + terms.score_loss.sum() + terms.adaptive_kl_loss.sum() + local_loss = local_total / context.valid_token_count + local_loss.backward() + + reduced_loss = local_loss.detach().clone() + reduced_gradient = local_log_probs.grad.detach().clone() + dist.all_reduce(reduced_loss, group=process_group) + dist.all_reduce(reduced_gradient, group=process_group) + torch.testing.assert_close(reduced_loss, oracle_loss.detach()) + torch.testing.assert_close(reduced_gradient, oracle_gradient) + + assert_partition([all_indices], singleton_groups[rank]) + assert_partition(list(torch.tensor_split(all_indices, 2))[rank % 2 : rank % 2 + 1], dp2_groups[rank // 2]) + assert_partition([torch.tensor_split(all_indices, world_size)[rank]], dist.group.WORLD) + + static_dp2_cp2 = [ + torch.tensor([0, 7, 8]), + torch.tensor([1, 6, 9]), + torch.tensor([2, 5, 10]), + torch.tensor([3, 4]), + ] + assert_partition([static_dp2_cp2[rank]], dist.group.WORLD) + + dynamic_cp = [ + [torch.tensor([0, 1]), torch.tensor([6])], + [torch.tensor([2]), torch.tensor([5, 7, 9])], + [torch.tensor([3, 4]), torch.tensor([8, 10])], + [torch.empty(0, dtype=torch.long)], + ] + assert_partition(dynamic_cp[rank], dist.group.WORLD) + finally: + dist.destroy_process_group() + + +def test_p3o_distributed_nonfinite_fails_synchronously(): + world_size = 2 + mp.spawn(_nonfinite_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) + + +def test_p3o_distributed_pipeline_broadcasts_last_stage_stats(): + world_size = 2 + mp.spawn(_pipeline_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) + + +def test_p3o_distributed_partition_and_objective_invariance(): + world_size = 4 + mp.spawn(_partition_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) diff --git a/tests/backends/megatron/test_p3o_loss.py b/tests/backends/megatron/test_p3o_loss.py new file mode 100644 index 000000000..ab0ad69f9 --- /dev/null +++ b/tests/backends/megatron/test_p3o_loss.py @@ -0,0 +1,173 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Metric-contract tests for the Megatron P3O loss branch. + +``relax.backends.megatron.loss`` imports ``megatron.core`` at module scope, and +CI installs no megatron. The branch under test only consumes token terms, so +the megatron surface is stubbed for the import and restored afterwards -- +keeping these assertions running in CI instead of silently skipping. +""" + +from argparse import Namespace + +import pytest +import torch + +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules + + +with stubbed_megatron_modules(("megatron", "ray")): + from relax.backends.megatron import loss as loss_module + +from relax.utils.training.p3o_utils import P3OStepContext + + +REQUIRED_P3O_METRICS = { + "p3o/normalized_ess", + "p3o/adaptive_cap", + "p3o/ratio_mean", + "p3o/ratio_std", + "p3o/cap_fraction", + "p3o/clip_fraction", + "p3o/score_loss", + "p3o/behavior_kl_proxy", + "p3o/adaptive_kl_loss", + "p3o/reference_kl", + "p3o/entropy", + "p3o/valid_tokens", + "p3o/total_loss", +} + + +def test_get_p3o_context_computes_micro_batch_scope_without_prepass(): + args = Namespace(p3o_ess_scope="micro-batch") + log_probs = torch.tensor([-0.4, -0.8]) + behavior_log_probs = torch.tensor([-0.5, -0.7]) + valid_mask = torch.tensor([True, True]) + + context = loss_module.get_p3o_context(args, log_probs, behavior_log_probs, valid_mask) + + assert context.valid_token_count.item() == 2 + assert 0.0 < context.normalized_ess.item() <= 1.0 + assert torch.equal(context.normalized_ess, context.adaptive_cap) + + +def test_get_p3o_context_rejects_unknown_scope(): + args = Namespace(p3o_ess_scope="window") + + with pytest.raises(ValueError, match="micro-batch.*step"): + loss_module.get_p3o_context(args, torch.zeros(1), torch.zeros(1), torch.ones(1, dtype=torch.bool)) + + +def test_p3o_loss_reports_complete_schema_without_reference_kl(monkeypatch): + step_context = P3OStepContext( + normalized_ess=torch.tensor(0.75, dtype=torch.float64), + adaptive_cap=torch.tensor(0.75, dtype=torch.float64), + valid_token_count=torch.tensor(2.0, dtype=torch.float64), + ratio_mean=torch.tensor(1.0, dtype=torch.float64), + ratio_std=torch.tensor(0.0, dtype=torch.float64), + ) + args = Namespace( + _p3o_step_context=step_context, + entropy_coef=0.0, + p3o_ess_scope="step", + qkv_format="thd", + use_kl_loss=False, + ) + log_probs = torch.tensor([-0.4, -0.8], requires_grad=True) + monkeypatch.setattr( + loss_module, + "get_log_probs_and_entropy", + lambda *args, **kwargs: ( + torch.empty(0), + { + "log_probs": [log_probs], + "entropy": [torch.tensor([0.2, 0.3])], + }, + ), + ) + monkeypatch.setattr( + loss_module, + "get_cp_local_valid_mask", + lambda *args, **kwargs: torch.tensor([True, True]), + ) + batch = { + "advantages": torch.tensor([1.0, -1.0]), + "rollout_log_probs": [log_probs.detach().clone()], + "unconcat_tokens": [torch.tensor([1, 2])], + "total_lengths": [2], + "response_lengths": [2], + "loss_masks": [torch.ones(2)], + } + + _, metrics = loss_module.p3o_loss_function(args, batch, torch.zeros(1), torch.sum) + + assert REQUIRED_P3O_METRICS <= metrics.keys() + assert not any(metric.startswith("opd/") for metric in metrics) + assert torch.equal(metrics["p3o/reference_kl"], torch.zeros(())) + assert not metrics["p3o/reference_kl"].requires_grad + + +def test_p3o_loss_function_normalizes_by_true_valid_tokens(monkeypatch): + """All-masked samples must not add phantom tokens to P3O's normalizer.""" + args = Namespace( + advantage_estimator="p3o", + allgather_cp=False, + calculate_per_token_loss=True, + global_batch_size=2, + loss_type="policy_loss", + qkv_format="thd", + recompute_loss_function=False, + use_opd=False, + ) + batch = { + "loss_masks": [torch.zeros(2), torch.tensor([1.0, 0.0])], + "response_lengths": [2, 2], + "total_lengths": [3, 3], + } + monkeypatch.setattr(loss_module, "get_cp_local_num_tokens", lambda *args, **kwargs: torch.tensor(2.0)) + monkeypatch.setattr(loss_module, "get_sum_of_sample_mean", lambda *args, **kwargs: torch.tensor(0.0)) + monkeypatch.setattr( + loss_module, + "get_cp_local_valid_mask", + lambda *args, **kwargs: torch.tensor([False, False, True, False]), + ) + monkeypatch.setattr( + loss_module, + "p3o_loss_function", + lambda *args, **kwargs: (torch.tensor(3.0, requires_grad=True), {"loss": torch.tensor(3.0)}), + ) + monkeypatch.setattr( + loss_module, + "policy_loss_function", + lambda *args, **kwargs: pytest.fail("P3O must not use the ordinary policy-loss path"), + ) + monkeypatch.setattr( + loss_module, + "compute_policy_opd_loss", + lambda *args, **kwargs: pytest.fail("P3O must not call compute_policy_opd_loss"), + ) + + _, normalizer, logging_dict = loss_module.loss_function(args, batch, 1, torch.zeros(1)) + + assert normalizer.item() == 1 + assert logging_dict["values"][0].item() == 1 + + +def test_policy_loss_dispatch_selects_dedicated_p3o_path(): + args = Namespace(advantage_estimator="p3o", use_opd=False) + + assert loss_module._select_policy_loss_function(args) is loss_module.p3o_loss_function + + +def test_policy_loss_dispatch_rejects_p3o_with_opd(): + args = Namespace(advantage_estimator="p3o", use_opd=True) + + with pytest.raises(ValueError, match="P3O and OPD are mutually exclusive"): + loss_module._select_policy_loss_function(args) + + +def test_policy_loss_dispatch_preserves_opd_for_non_p3o_estimators(): + args = Namespace(advantage_estimator="grpo", use_opd=True) + + assert loss_module._select_policy_loss_function(args) is loss_module.policy_loss_function diff --git a/tests/backends/megatron/test_p3o_model_step.py b/tests/backends/megatron/test_p3o_model_step.py new file mode 100644 index 000000000..4499cbfea --- /dev/null +++ b/tests/backends/megatron/test_p3o_model_step.py @@ -0,0 +1,69 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Exception-safety tests for the P3O optimizer-step lifecycle.""" + +from __future__ import annotations + +import ast +import sys +from argparse import Namespace +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import patch + +import pytest + +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules + + +MODEL_PATH = Path(__file__).resolve().parents[3] / "relax" / "backends" / "megatron" / "model.py" + +stream_dataloader = ModuleType("relax.utils.data.stream_dataloader") +stream_dataloader.StreamingTQIterator = object + +with ( + patch.dict(sys.modules, {"relax.utils.data.stream_dataloader": stream_dataloader}), + stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), +): + from relax.backends.megatron.model import _preserved_dynamic_cp_group + + +def test_p3o_model_step_restores_dynamic_cp_group_after_error(): + original_group = object() + dynamic_group = object() + inner = SimpleNamespace(pg_collection=SimpleNamespace(cp=original_group)) + wrapped = SimpleNamespace(module=inner) + args = Namespace(dynamic_context_parallel=True) + + with pytest.raises(RuntimeError, match="stats pass failed"): + with _preserved_dynamic_cp_group(args, [wrapped]): + inner.pg_collection.cp = dynamic_group + raise RuntimeError("stats pass failed") + + assert inner.pg_collection.cp is original_group + + +def test_p3o_model_step_guard_covers_stats_and_train_passes(): + tree = ast.parse(MODEL_PATH.read_text(encoding="utf-8")) + train_one_step = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "train_one_step" + ) + guard = next( + node + for node in ast.walk(train_one_step) + if isinstance(node, ast.With) + and any( + isinstance(child, ast.Name) and child.id == "_preserved_dynamic_cp_group" + for item in node.items + for child in ast.walk(item.context_expr) + ) + ) + guarded_calls = { + child.func.id for child in ast.walk(guard) if isinstance(child, ast.Call) and isinstance(child.func, ast.Name) + } + assert "compute_p3o_step_context" in guarded_calls + assert "forward_backward_func" in guarded_calls + guarded_source = ast.dump(guard) + assert "p3o_ess_scope" in guarded_source + assert "micro-batch" in guarded_source + assert "step" in guarded_source diff --git a/tests/backends/megatron/test_p3o_observability.py b/tests/backends/megatron/test_p3o_observability.py new file mode 100644 index 000000000..0c9d50304 --- /dev/null +++ b/tests/backends/megatron/test_p3o_observability.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Behavior tests for P3O rollout-policy age observability.""" + +import ast +from pathlib import Path + +import pytest + +from relax.backends.megatron.rollout_policy_lag import ( + ROLLOUT_POLICY_TAG, + build_rollout_policy_age_metrics, + compute_rollout_policy_age_rollouts, + initial_rollout_policy_snapshot_rollout, + maybe_refresh_rollout_policy, + rollout_weights_tag, + should_refresh_rollout_policy, + validate_update_weights_interval, +) + + +class _RecordingBackuper: + def __init__(self) -> None: + self.copies: list[tuple[str, str]] = [] + + def copy(self, *, src_tag: str, dst_tag: str) -> None: + self.copies.append((src_tag, dst_tag)) + + +@pytest.mark.parametrize("interval", [0, -1, -10]) +def test_rollout_policy_interval_rejects_invalid_values(interval): + with pytest.raises(ValueError, match="positive integer"): + validate_update_weights_interval(interval) + + +def test_rollout_policy_age_uses_rollout_units(): + assert compute_rollout_policy_age_rollouts(15, 11) == 4 + + +@pytest.mark.parametrize( + ("current_rollout_id", "snapshot_rollout_id", "message"), + [ + (-1, 0, "current_rollout_id"), + (0, -1, "snapshot_rollout_id"), + (2, 3, "cannot precede"), + ], +) +def test_rollout_policy_age_rejects_invalid_versions(current_rollout_id, snapshot_rollout_id, message): + with pytest.raises(ValueError, match=message): + compute_rollout_policy_age_rollouts(current_rollout_id, snapshot_rollout_id) + + +def test_rollout_policy_age_interval_three_sequence(): + snapshot_rollout = 0 + observed = [] + refreshes = [] + backuper = _RecordingBackuper() + + for rollout_id in range(6): + observed.append(compute_rollout_policy_age_rollouts(rollout_id, snapshot_rollout)) + refreshed = maybe_refresh_rollout_policy(backuper, rollout_id, 3, 7) + refreshes.append(refreshed) + if refreshed: + snapshot_rollout = rollout_id + 1 + + assert observed == [0, 1, 2, 0, 1, 2] + assert refreshes == [False, False, True, False, False, True] + assert backuper.copies == [("actor", ROLLOUT_POLICY_TAG), ("actor", ROLLOUT_POLICY_TAG)] + + +def test_rollout_policy_snapshot_initializes_for_fresh_and_resumed_runs(): + assert initial_rollout_policy_snapshot_rollout(0) == 0 + assert initial_rollout_policy_snapshot_rollout(101) == 101 + assert compute_rollout_policy_age_rollouts(101, initial_rollout_policy_snapshot_rollout(101)) == 0 + + +def test_rollout_policy_snapshot_rejects_invalid_resume_version(): + with pytest.raises(ValueError, match="start_rollout_id"): + initial_rollout_policy_snapshot_rollout(-1) + + +def test_rollout_policy_age_metrics_have_exact_keys_and_values(): + assert build_rollout_policy_age_metrics(current_rollout_id=7, rollout_policy_snapshot_rollout=5) == { + "train/current_rollout_id": 7, + "train/rollout_policy_snapshot_rollout": 5, + "train/p3o/rollout_policy_age_rollouts": 2, + } + + +def test_rollout_policy_refresh_calls_backuper_only_at_boundary(): + backuper = _RecordingBackuper() + + assert not maybe_refresh_rollout_policy(backuper, rollout_id=0, update_weights_interval=3, num_rollout=6) + assert backuper.copies == [] + + assert maybe_refresh_rollout_policy(backuper, rollout_id=2, update_weights_interval=3, num_rollout=6) + assert backuper.copies == [("actor", ROLLOUT_POLICY_TAG)] + + +def test_on_policy_mode_uses_actor_and_refreshes_every_rollout(): + assert rollout_weights_tag(1) == "actor" + assert should_refresh_rollout_policy(5, 1, 10) + + +def test_periodic_sync_mode_uses_rollout_policy_snapshot(): + assert rollout_weights_tag(3) == ROLLOUT_POLICY_TAG + + +def test_final_rollout_forces_refresh_away_from_interval_boundary(): + backuper = _RecordingBackuper() + + assert maybe_refresh_rollout_policy(backuper, rollout_id=4, update_weights_interval=3, num_rollout=5) + assert backuper.copies == [("actor", ROLLOUT_POLICY_TAG)] + + +def test_hybrid_training_publishes_snapshot_rollout_before_train(): + actor_path = Path(__file__).resolve().parents[3] / "relax" / "backends" / "megatron" / "actor.py" + tree = ast.parse(actor_path.read_text(encoding="utf-8")) + actor_class = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "MegatronTrainRayActor" + ) + train_hybrid = next( + node for node in actor_class.body if isinstance(node, ast.FunctionDef) and node.name == "train_hybrid" + ) + + snapshot_assignment = next( + node + for node in ast.walk(train_hybrid) + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Attribute) and target.attr == "rollout_policy_snapshot_rollout" + for target in node.targets + ) + ) + train_call = next( + node + for node in ast.walk(train_hybrid) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "train" + ) + + assert snapshot_assignment.lineno < train_call.lineno diff --git a/tests/backends/megatron/test_p3o_on_policy.py b/tests/backends/megatron/test_p3o_on_policy.py new file mode 100644 index 000000000..c7d02e215 --- /dev/null +++ b/tests/backends/megatron/test_p3o_on_policy.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tiny-model acceptance gate for P3O's on-policy degeneration.""" + +import copy + +import torch +from torch import nn + +from relax.utils.training.p3o_utils import ( + compute_p3o_sufficient_stats, + compute_p3o_token_terms, + finalize_p3o_step_context, +) + + +def _flatten_gradients(model: nn.Module) -> torch.Tensor: + return torch.cat([parameter.grad.flatten() for parameter in model.parameters()]) + + +def test_p3o_on_policy_matches_policy_gradient_and_parameter_update(): + torch.manual_seed(42) + base_model = nn.Linear(3, 1, bias=True) + pg_model = copy.deepcopy(base_model) + p3o_model = copy.deepcopy(base_model) + features = torch.tensor( + [ + [0.2, -0.5, 1.0], + [1.5, 0.3, -0.7], + [-0.4, 0.8, 0.1], + [0.9, -1.2, 0.6], + [-0.8, -0.2, 1.3], + [0.5, 0.7, -0.9], + ], + dtype=torch.float32, + ) + advantages = torch.tensor([1.0, -0.5, 0.75, -1.25, 0.4, 0.9]) + valid_mask = torch.ones(features.size(0), dtype=torch.bool) + behavior_log_probs = base_model(features).squeeze(-1).detach() + + pg_optimizer = torch.optim.SGD(pg_model.parameters(), lr=0.05) + pg_log_probs = pg_model(features).squeeze(-1) + pg_loss = -(pg_log_probs * advantages).mean() + pg_loss.backward() + pg_gradients = _flatten_gradients(pg_model).clone() + + p3o_optimizer = torch.optim.SGD(p3o_model.parameters(), lr=0.05) + p3o_log_probs = p3o_model(features).squeeze(-1) + context = finalize_p3o_step_context(compute_p3o_sufficient_stats(p3o_log_probs, behavior_log_probs, valid_mask)) + terms = compute_p3o_token_terms( + p3o_log_probs, + behavior_log_probs, + advantages, + valid_mask, + context, + ) + p3o_loss = (terms.score_loss + terms.adaptive_kl_loss).mean() + p3o_loss.backward() + p3o_gradients = _flatten_gradients(p3o_model).clone() + + cosine = torch.nn.functional.cosine_similarity(pg_gradients, p3o_gradients, dim=0) + relative_l2 = torch.linalg.vector_norm(p3o_gradients - pg_gradients) / torch.linalg.vector_norm(pg_gradients) + assert float(cosine) >= 0.9999 + assert float(relative_l2) <= 1e-4 + assert float(terms.adaptive_kl_loss.detach().abs().max()) <= 1e-7 + + pg_optimizer.step() + p3o_optimizer.step() + for pg_parameter, p3o_parameter in zip(pg_model.parameters(), p3o_model.parameters(), strict=True): + torch.testing.assert_close(p3o_parameter, pg_parameter, rtol=1e-4, atol=1e-6) diff --git a/tests/backends/megatron/test_p3o_partition_invariance.py b/tests/backends/megatron/test_p3o_partition_invariance.py new file mode 100644 index 000000000..0eda5033b --- /dev/null +++ b/tests/backends/megatron/test_p3o_partition_invariance.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Logical partition-invariance tests for optimizer-step P3O.""" + +import math + +import pytest +import torch + +from relax.utils.training.p3o_utils import ( + P3OSufficientStats, + compute_p3o_sufficient_stats, + compute_p3o_token_terms, + finalize_p3o_step_context, +) + + +TOKEN_COUNT = 11 +INDICES = torch.arange(TOKEN_COUNT) +BEHAVIOR_LOG_PROBS = torch.full((TOKEN_COUNT,), -2.0) +LOG_RATIOS = torch.tensor([math.log(value) for value in (1.0, 2.0, 0.5, 4.0, 0.8, 1.4, 0.25, 3.0, 1.1, 0.6, 2.5)]) +ADVANTAGES = torch.tensor([1.0, -1.0, 0.5, 2.0, -0.2, 0.7, -1.5, 1.2, 0.4, -0.8, 1.8]) +VALID_MASK = torch.tensor([True, True, False, True, True, False, True, True, True, False, True]) + + +def _evaluate(shards: list[torch.Tensor]): + log_probs = (BEHAVIOR_LOG_PROBS + LOG_RATIOS).clone().requires_grad_(True) + stats = P3OSufficientStats.zeros() + for shard in shards: + if shard.numel() == 0: + stats = stats + P3OSufficientStats.zeros() + continue + stats = stats + compute_p3o_sufficient_stats( + log_probs[shard], + BEHAVIOR_LOG_PROBS[shard], + VALID_MASK[shard], + ) + context = finalize_p3o_step_context(stats) + + total = 0.0 * log_probs.sum() + for shard in shards: + if shard.numel() == 0: + continue + terms = compute_p3o_token_terms( + log_probs[shard], + BEHAVIOR_LOG_PROBS[shard], + ADVANTAGES[shard], + VALID_MASK[shard], + context, + ) + total = total + terms.score_loss.sum() + terms.adaptive_kl_loss.sum() + loss = total / context.valid_token_count + loss.backward() + return context, loss.detach(), log_probs.grad.detach() + + +def _assert_matches_oracle(shards: list[torch.Tensor]): + expected_context, expected_loss, expected_grad = _evaluate([INDICES]) + actual_context, actual_loss, actual_grad = _evaluate(shards) + + torch.testing.assert_close(actual_context.normalized_ess, expected_context.normalized_ess) + torch.testing.assert_close(actual_context.adaptive_cap, expected_context.adaptive_cap) + torch.testing.assert_close(actual_context.ratio_mean, expected_context.ratio_mean) + torch.testing.assert_close(actual_context.ratio_std, expected_context.ratio_std) + torch.testing.assert_close(actual_context.valid_token_count, expected_context.valid_token_count) + torch.testing.assert_close(actual_loss, expected_loss) + torch.testing.assert_close(actual_grad, expected_grad) + + +@pytest.mark.parametrize("micro_batch_size", [1, 2, 4]) +def test_p3o_partition_invariance_fixed_micro_batches(micro_batch_size): + shards = list(torch.split(INDICES, micro_batch_size)) + _assert_matches_oracle(shards) + + +def test_p3o_partition_invariance_ragged_and_dummy_micro_batches(): + shards = [ + INDICES[0:3], + torch.empty(0, dtype=torch.long), + INDICES[3:4], + INDICES[4:9], + torch.empty(0, dtype=torch.long), + INDICES[9:], + ] + _assert_matches_oracle(shards) + + +@pytest.mark.parametrize("data_parallel_size", [1, 2, 4]) +def test_p3o_partition_invariance_logical_data_parallel_shards(data_parallel_size): + _assert_matches_oracle(list(torch.tensor_split(INDICES, data_parallel_size))) + + +def test_p3o_partition_invariance_static_dp2_cp2_zigzag_shards(): + shards = [ + torch.tensor([0, 7, 8]), + torch.tensor([1, 6, 9]), + torch.tensor([2, 5, 10]), + torch.tensor([3, 4]), + ] + _assert_matches_oracle(shards) + + +def test_p3o_partition_invariance_dynamic_cp_and_zero_local_tokens(): + shards = [ + torch.tensor([0, 1, 6]), + torch.tensor([2, 5, 7, 9]), + torch.tensor([3, 4, 8, 10]), + torch.empty(0, dtype=torch.long), + ] + _assert_matches_oracle(shards) diff --git a/tests/backends/megatron/test_p3o_step.py b/tests/backends/megatron/test_p3o_step.py new file mode 100644 index 000000000..b131a189a --- /dev/null +++ b/tests/backends/megatron/test_p3o_step.py @@ -0,0 +1,527 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for optimizer-step P3O stats synchronization.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules + + +with stubbed_megatron_modules(("megatron", "ray", "tensordict")): + from relax.backends.megatron import cp_utils, p3o_step + from relax.backends.megatron.p3o_step import synchronize_p3o_stats + +from relax.utils.training.p3o_utils import P3OSufficientStats + + +@pytest.fixture(autouse=True) +def _stub_cp_world_size(monkeypatch): + monkeypatch.setattr( + cp_utils, + "mpu", + SimpleNamespace(get_context_parallel_world_size=lambda: 1), + ) + # The target SIF has a real Megatron installation, whereas lightweight + # developer environments use the module stubs above. Keep these unit tests + # hermetic in both cases instead of querying an uninitialized PP group. + monkeypatch.setattr(p3o_step.mpu, "is_pipeline_last_stage", lambda ignore_virtual=False: True) + + +def _stats(values: tuple[float, float, float]) -> P3OSufficientStats: + vector = torch.tensor(values, dtype=torch.float64) + return P3OSufficientStats.from_vector(vector) + + +def test_p3o_step_single_pipeline_stage_preserves_stats(monkeypatch): + monkeypatch.setattr(torch.distributed, "is_available", lambda: False) + stats = _stats((7.5, 21.25, 4.0)) + + synchronized = synchronize_p3o_stats( + stats, + torch.zeros((), dtype=torch.float64), + dp_cp_group=None, + pp_group=None, + is_pipeline_last_stage=True, + ) + + torch.testing.assert_close(synchronized.as_vector(), stats.as_vector(), rtol=0.0, atol=0.0) + + +def test_p3o_step_non_last_stage_receives_pipeline_last_stats(monkeypatch): + expected = torch.tensor([7.5, 21.25, 4.0, 0.0], dtype=torch.float64) + pp_group = object() + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group: 2) + + def fail_if_reduced(*args, **kwargs): + raise AssertionError("a non-last PP stage must not reduce token stats over DP x CP") + + def broadcast_from_last(vector, *, group, group_src): + assert group is pp_group + assert group_src == 1 + vector.copy_(expected) + + monkeypatch.setattr(torch.distributed, "all_reduce", fail_if_reduced) + monkeypatch.setattr(torch.distributed, "broadcast", broadcast_from_last) + + synchronized = synchronize_p3o_stats( + P3OSufficientStats.zeros(), + torch.zeros((), dtype=torch.float64), + dp_cp_group=None, + pp_group=pp_group, + is_pipeline_last_stage=False, + ) + + torch.testing.assert_close(synchronized.as_vector(), expected[:3], rtol=0.0, atol=0.0) + + +def test_p3o_step_raises_only_after_global_invalid_flag_is_visible(monkeypatch): + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + dp_cp_group = object() + + def all_reduce(vector, *, op, group): + vector[3] = 1.0 + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + with pytest.raises(ValueError, match="non-finite importance ratio"): + synchronize_p3o_stats( + _stats((1.0, 1.0, 1.0)), + torch.zeros((), dtype=torch.float64), + dp_cp_group=dp_cp_group, + pp_group=None, + is_pipeline_last_stage=True, + ) + + +def test_compute_p3o_step_context_plain_text_forward_kwargs(monkeypatch): + """ESS pre-pass forward_step must use tokens+packed_seq_params for plain + text.""" + from argparse import Namespace + + captured = {} + + def fake_model(**kwargs): + captured.update(kwargs) + return torch.zeros(1, 1, 768) + + def fake_get_batch(iterator, keys, *_args, **_kwargs): + return { + "tokens": torch.zeros(4, dtype=torch.long), + "packed_seq_params": "packed_sentinel", + "total_lengths": [4], + "response_lengths": [2], + "loss_masks": [torch.ones(4)], + "rollout_log_probs": [torch.zeros(4)], + "full_loss_masks": torch.ones(4), + "unconcat_tokens": [torch.zeros(4, dtype=torch.long)], + } + + def fake_forward_backward(forward_step_func, data_iterator, model, **_kwargs): + # Call forward_step once to trigger kwarg capture (avoid calling collect callback) + forward_step_func(data_iterator[0], model[0]) + return None + + # Prevent the lazy `from .loss import get_log_probs_and_entropy` from executing + # by ensuring the forward_backward func never calls the collect callback + monkeypatch.setattr(p3o_step, "get_batch", fake_get_batch) + monkeypatch.setattr(p3o_step, "get_forward_backward_func", lambda: fake_forward_backward) + monkeypatch.setattr(p3o_step, "synchronize_p3o_stats", lambda *_, **__: _stats((7.5, 21.25, 4.0))) + monkeypatch.setattr( + p3o_step, + "finalize_p3o_step_context", + lambda s: p3o_step.P3OStepContext( + normalized_ess=torch.tensor(0.66), + adaptive_cap=torch.tensor(0.66), + valid_token_count=torch.tensor(4.0), + ratio_mean=torch.tensor(1.875), + ratio_std=torch.tensor(0.5), + clamp_events=0, + ), + ) + monkeypatch.setattr(torch, "no_grad", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_iterator_positions", lambda _: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_rng_state", lambda: __import__("contextlib").nullcontext()) + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=None, + ) + # loss.py is a lazy import inside compute_p3o_step_context (line 140 of p3o_step.py). + # It fires after the stubbed_megatron_modules context has already exited, so we must + # inject a mock for loss before the function is called. + monkeypatch.setitem(sys.modules, "relax.backends.megatron.loss", MagicMock()) + p3o_step.compute_p3o_step_context(args, [iter([None])], [fake_model], num_microbatches=1) + + assert captured["input_ids"] is not None + assert str(captured["input_ids"].dtype) == "torch.int64" + assert captured["packed_seq_params"] == "packed_sentinel" + assert captured["loss_mask"] is not None + + +def test_compute_p3o_step_context_matches_training_multimodal_kwarg_gate(monkeypatch): + """ESS pre-pass must not pass multimodal kwargs that training omits.""" + from argparse import Namespace + + captured = {} + + def fake_model(**kwargs): + captured.update(kwargs) + return torch.zeros(1, 1, 768) + + def fake_get_batch(iterator, keys, *_args, **_kwargs): + return { + "tokens": torch.zeros(4, dtype=torch.long), + "packed_seq_params": "packed_sentinel", + "multimodal_train_inputs": {"pixel_values": torch.ones(1)}, + "total_lengths": [4], + "response_lengths": [2], + "loss_masks": [torch.ones(4)], + "rollout_log_probs": [torch.zeros(4)], + "full_loss_masks": torch.ones(4), + "unconcat_tokens": [torch.zeros(4, dtype=torch.long)], + } + + def fake_forward_backward(forward_step_func, data_iterator, model, **_kwargs): + forward_step_func(data_iterator[0], model[0]) + return None + + monkeypatch.setattr(p3o_step, "get_batch", fake_get_batch) + monkeypatch.setattr(p3o_step, "get_forward_backward_func", lambda: fake_forward_backward) + monkeypatch.setattr(p3o_step, "synchronize_p3o_stats", lambda *_, **__: _stats((7.5, 21.25, 4.0))) + monkeypatch.setattr( + p3o_step, + "finalize_p3o_step_context", + lambda _: p3o_step.P3OStepContext( + normalized_ess=torch.tensor(0.66), + adaptive_cap=torch.tensor(0.66), + valid_token_count=torch.tensor(4.0), + ratio_mean=torch.tensor(1.875), + ratio_std=torch.tensor(0.5), + ), + ) + monkeypatch.setattr(torch, "no_grad", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_iterator_positions", lambda _: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_rng_state", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setitem(sys.modules, "relax.backends.megatron.loss", MagicMock()) + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + is_vl_model=False, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=None, + ) + p3o_step.compute_p3o_step_context(args, [iter([None])], [fake_model], num_microbatches=1) + + assert "pixel_values" not in captured + + +def test_compute_p3o_step_context_vl_unsplit_forward_kwargs(monkeypatch): + """ESS pre-pass forward_step must use unsplit_tokens for VL models.""" + from argparse import Namespace + + captured = {} + + def fake_model(**kwargs): + captured.update(kwargs) + return torch.zeros(1, 1, 768) + + def fake_get_batch(iterator, keys, *_args, **_kwargs): + return { + "tokens": torch.zeros(4, dtype=torch.long), + "unsplit_tokens": torch.zeros(8, dtype=torch.long), # VL path + "packed_seq_params": "packed_sentinel", + "total_lengths": [4], + "response_lengths": [2], + "loss_masks": [torch.ones(4)], + "rollout_log_probs": [torch.zeros(4)], + "full_loss_masks": torch.ones(4), + "unconcat_tokens": [torch.zeros(4, dtype=torch.long)], + } + + def fake_forward_backward(forward_step_func, data_iterator, model, **_kwargs): + output_tensor, _ = forward_step_func(data_iterator[0], model[0]) + return None + + monkeypatch.setattr(p3o_step, "get_batch", fake_get_batch) + monkeypatch.setattr(p3o_step, "get_forward_backward_func", lambda: fake_forward_backward) + monkeypatch.setattr(p3o_step, "synchronize_p3o_stats", lambda *_, **__: _stats((7.5, 21.25, 4.0))) + monkeypatch.setattr( + p3o_step, + "finalize_p3o_step_context", + lambda s: p3o_step.P3OStepContext( + normalized_ess=torch.tensor(0.66), + adaptive_cap=torch.tensor(0.66), + valid_token_count=torch.tensor(4.0), + ratio_mean=torch.tensor(1.875), + ratio_std=torch.tensor(0.5), + clamp_events=0, + ), + ) + monkeypatch.setattr(torch, "no_grad", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_iterator_positions", lambda _: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_rng_state", lambda: __import__("contextlib").nullcontext()) + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + is_vl_model=True, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=None, + ) + # cp_utils.maybe_padded_total_lengths queries mpu for the CP world size; this + # test is single-process, so report CP=1 instead of a bare MagicMock. + monkeypatch.setattr(cp_utils.mpu, "get_context_parallel_world_size", lambda: 1) + monkeypatch.setitem(sys.modules, "relax.backends.megatron.loss", MagicMock()) + p3o_step.compute_p3o_step_context(args, [iter([None])], [fake_model], num_microbatches=1) + + # VL path: should use unsplit_tokens, packed_seq_params=None + assert captured["input_ids"].shape == (8,), "VL path must use unsplit_tokens" + assert captured["packed_seq_params"] is None, "VL path sets packed_seq_params=None" + assert captured["loss_mask"] is not None + + +def test_compute_p3o_step_context_vl_thd_bridge_forward_kwargs(monkeypatch): + """ESS pre-pass forward_step must use thd bridge path + (vlm_packed_seq_params, loss_mask=None).""" + from argparse import Namespace + + captured = {} + + def fake_model(**kwargs): + captured.update(kwargs) + return torch.zeros(1, 1, 768) + + def fake_get_batch(iterator, keys, *_args, **_kwargs): + return { + "tokens": torch.zeros(4, dtype=torch.long), + "unsplit_tokens": torch.zeros(8, dtype=torch.long), + "unsplit_attention_mask": torch.ones(8), + "vlm_packed_seq_params": "vlm_packed_sentinel", # thd bridge marker + "packed_seq_params": "packed_sentinel", + "total_lengths": [4], + "response_lengths": [2], + "loss_masks": [torch.ones(4)], + "rollout_log_probs": [torch.zeros(4)], + "full_loss_masks": torch.ones(4), + "unconcat_tokens": [torch.zeros(4, dtype=torch.long)], + } + + def fake_forward_backward(forward_step_func, data_iterator, model, **_kwargs): + output_tensor, _ = forward_step_func(data_iterator[0], model[0]) + return None + + monkeypatch.setattr(p3o_step, "get_batch", fake_get_batch) + monkeypatch.setattr(p3o_step, "get_forward_backward_func", lambda: fake_forward_backward) + monkeypatch.setattr(p3o_step, "synchronize_p3o_stats", lambda *_, **__: _stats((7.5, 21.25, 4.0))) + monkeypatch.setattr( + p3o_step, + "finalize_p3o_step_context", + lambda s: p3o_step.P3OStepContext( + normalized_ess=torch.tensor(0.66), + adaptive_cap=torch.tensor(0.66), + valid_token_count=torch.tensor(4.0), + ratio_mean=torch.tensor(1.875), + ratio_std=torch.tensor(0.5), + clamp_events=0, + ), + ) + monkeypatch.setattr(torch, "no_grad", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_iterator_positions", lambda _: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_rng_state", lambda: __import__("contextlib").nullcontext()) + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + is_vl_model=True, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=None, + ) + monkeypatch.setattr(cp_utils.mpu, "get_context_parallel_world_size", lambda: 1) + monkeypatch.setitem(sys.modules, "relax.backends.megatron.loss", MagicMock()) + p3o_step.compute_p3o_step_context(args, [iter([None])], [fake_model], num_microbatches=1) + + # thd bridge path: unsplit_tokens, vlm_packed_seq_params, unsplit_attention_mask, loss_mask=None + assert captured["input_ids"].shape == (8,), "thd bridge must use unsplit_tokens" + assert captured["packed_seq_params"] == "vlm_packed_sentinel", "thd bridge uses vlm_packed_seq_params" + assert captured["attention_mask"] is not None, "thd bridge requires attention_mask" + assert captured["loss_mask"] is None, "thd bridge sets loss_mask=None" + + +def test_compute_p3o_step_context_dynamic_cp_group_switching(monkeypatch): + """ESS pre-pass forward_step must switch pg_collection.cp for dynamic + CP.""" + from argparse import Namespace + + captured_pg = [] + orig_cp_group = object() + dynamic_cp_group = object() + + class FakePGCollection: + def __init__(self): + self.cp = orig_cp_group + + class FakeInner: + def __init__(self): + self.pg_collection = FakePGCollection() + + class FakeModel: + def __init__(self): + self.module = FakeInner() + + def __call__(self, **kwargs): + captured_pg.append(self.module.pg_collection.cp) + return torch.zeros(1, 1, 768) + + fake_model = FakeModel() + + batch = { + "tokens": torch.zeros(4, dtype=torch.long), + "unsplit_tokens": torch.zeros(8, dtype=torch.long), + "packed_seq_params": "packed_sentinel", + "dynamic_cp_size": 2, # trigger dynamic CP path + "padded_total_lengths": [8], + "total_lengths": [4], + "response_lengths": [2], + "loss_masks": [torch.ones(4)], + "rollout_log_probs": [torch.zeros(4)], + "full_loss_masks": torch.ones(4), + "unconcat_tokens": [torch.zeros(4, dtype=torch.long)], + } + + def fake_get_batch(iterator, keys, *_args, **_kwargs): + return batch + + def fake_forward_backward(forward_step_func, data_iterator, model, **_kwargs): + output_tensor, _ = forward_step_func(data_iterator[0], model[0]) + return None + + monkeypatch.setattr(p3o_step.mpu, "get_dynamic_data_context_parallel_groups", lambda group_size: dynamic_cp_group) + monkeypatch.setattr(p3o_step, "get_batch", fake_get_batch) + monkeypatch.setattr(p3o_step, "get_forward_backward_func", lambda: fake_forward_backward) + monkeypatch.setattr(p3o_step, "synchronize_p3o_stats", lambda *_, **__: _stats((7.5, 21.25, 4.0))) + monkeypatch.setattr( + p3o_step, + "finalize_p3o_step_context", + lambda s: p3o_step.P3OStepContext( + normalized_ess=torch.tensor(0.66), + adaptive_cap=torch.tensor(0.66), + valid_token_count=torch.tensor(4.0), + ratio_mean=torch.tensor(1.875), + ratio_std=torch.tensor(0.5), + clamp_events=0, + ), + ) + monkeypatch.setattr(torch, "no_grad", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_iterator_positions", lambda _: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_rng_state", lambda: __import__("contextlib").nullcontext()) + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + is_vl_model=True, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=None, + ) + monkeypatch.setattr(cp_utils.mpu, "get_context_parallel_world_size", lambda: 1) + monkeypatch.setitem(sys.modules, "relax.backends.megatron.loss", MagicMock()) + p3o_step.compute_p3o_step_context(args, [iter([None])], [fake_model], num_microbatches=1) + + # The forward should have been called with dynamic_cp_group active + assert len(captured_pg) == 1, "forward_step should call model_chunk once" + assert captured_pg[0] is dynamic_cp_group, "pg_collection.cp must switch to dynamic group during forward" + # After forward, it should be restored (verify via the finally block's side effect) + assert fake_model.module.pg_collection.cp is orig_cp_group, "pg_collection.cp must be restored after forward" + assert batch["padded_total_lengths"] == [8], "dynamic-CP padding metadata must not be overwritten" + + +def test_compute_p3o_step_context_dynamic_cp_one_does_not_add_static_padding(monkeypatch): + """Dynamic CP size one must not inherit padding from the static CP + group.""" + from argparse import Namespace + + batch = { + "tokens": torch.zeros(4, dtype=torch.long), + "unsplit_tokens": torch.zeros(4, dtype=torch.long), + "packed_seq_params": "packed_sentinel", + "dynamic_cp_size": 1, + "total_lengths": [4], + "response_lengths": [2], + "loss_masks": [torch.ones(4)], + "rollout_log_probs": [torch.zeros(4)], + "full_loss_masks": torch.ones(4), + "unconcat_tokens": [torch.zeros(4, dtype=torch.long)], + } + + class FakePGCollection: + cp = object() + + class FakeInner: + pg_collection = FakePGCollection() + + class FakeModel: + module = FakeInner() + + def __call__(self, **kwargs): + return torch.zeros(1, 1, 768) + + def fake_forward_backward(forward_step_func, data_iterator, model, **_kwargs): + forward_step_func(data_iterator[0], model[0]) + return None + + monkeypatch.setattr(p3o_step, "get_batch", lambda *args, **kwargs: batch) + monkeypatch.setattr(p3o_step, "get_forward_backward_func", lambda: fake_forward_backward) + monkeypatch.setattr(p3o_step.mpu, "get_dynamic_data_context_parallel_groups", lambda group_size: object()) + monkeypatch.setattr(p3o_step, "synchronize_p3o_stats", lambda *_, **__: _stats((7.5, 21.25, 4.0))) + monkeypatch.setattr( + p3o_step, + "finalize_p3o_step_context", + lambda _: p3o_step.P3OStepContext( + normalized_ess=torch.tensor(0.66), + adaptive_cap=torch.tensor(0.66), + valid_token_count=torch.tensor(4.0), + ratio_mean=torch.tensor(1.875), + ratio_std=torch.tensor(0.5), + ), + ) + monkeypatch.setattr(torch, "no_grad", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_iterator_positions", lambda _: __import__("contextlib").nullcontext()) + monkeypatch.setattr(p3o_step, "preserved_rng_state", lambda: __import__("contextlib").nullcontext()) + monkeypatch.setitem(sys.modules, "relax.backends.megatron.loss", MagicMock()) + monkeypatch.setattr(cp_utils.mpu, "get_context_parallel_world_size", lambda: 4) + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + is_vl_model=True, + seq_length=512, + micro_batch_size=1, + decoder_seq_length=None, + ) + p3o_step.compute_p3o_step_context(args, [iter([None])], [FakeModel()], num_microbatches=1) + + assert "padded_total_lengths" not in batch diff --git a/tests/backends/megatron/test_rollout_policy_lag.py b/tests/backends/megatron/test_rollout_policy_lag.py new file mode 100644 index 000000000..a77df9f22 --- /dev/null +++ b/tests/backends/megatron/test_rollout_policy_lag.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for periodic rollout policy snapshot scheduling.""" + +import pytest + +from relax.backends.megatron.rollout_policy_lag import ( + ROLLOUT_POLICY_TAG, + maybe_refresh_rollout_policy, + rollout_weights_tag, + should_refresh_rollout_policy, + validate_update_weights_interval, +) + + +class _RecordingBackuper: + def __init__(self): + self.copies = [] + + def copy(self, *, src_tag: str, dst_tag: str) -> None: + self.copies.append((src_tag, dst_tag)) + + +def test_rollout_policy_lag_interval_one_preserves_actor_updates(): + assert rollout_weights_tag(1) == "actor" + assert all(should_refresh_rollout_policy(step, 1, 5) for step in range(5)) + + +def test_rollout_policy_lag_interval_three_refreshes_boundaries_and_final_step(): + refreshes = [should_refresh_rollout_policy(step, 3, 8) for step in range(8)] + + assert rollout_weights_tag(3) == ROLLOUT_POLICY_TAG + assert refreshes == [False, False, True, False, False, True, False, True] + + +def test_rollout_policy_lag_copies_only_at_scheduled_boundaries(): + backuper = _RecordingBackuper() + + refreshed = [maybe_refresh_rollout_policy(backuper, step, 3, 8) for step in range(8)] + + assert refreshed == [False, False, True, False, False, True, False, True] + assert backuper.copies == [("actor", ROLLOUT_POLICY_TAG)] * 3 + + +@pytest.mark.parametrize("interval", [0, -1]) +def test_rollout_policy_lag_rejects_non_positive_intervals(interval): + with pytest.raises(ValueError, match="positive integer"): + validate_update_weights_interval(interval) diff --git a/tests/components/test_p3o_advantages.py b/tests/components/test_p3o_advantages.py new file mode 100644 index 000000000..f5a9559fc --- /dev/null +++ b/tests/components/test_p3o_advantages.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""P3O advantage-path parity with GRPO.""" + +import sys +from pathlib import Path +from types import SimpleNamespace + +import torch + + +# `relax.components.advantages` imports `megatron.core` at module level. CI installs no +# megatron, so the import runs under the shared stub; the advantage path under test is +# pure PyTorch and touches no megatron symbol at call time. +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "backends" / "megatron")) + +from _megatron_stub import stubbed_megatron_modules # noqa: E402 + + +with stubbed_megatron_modules(): + from relax.components.advantages import Advantages # noqa: E402 + + +def _compute(estimator: str): + advantages_class = Advantages.func_or_class + component = advantages_class.__new__(advantages_class) + component.config = SimpleNamespace( + advantage_estimator=estimator, + kl_coef=0.0, + use_kl_loss=False, + use_rollout_logprobs=True, + use_opd=False, + ) + rollout_data = { + "rollout_log_probs": [ + torch.tensor([-0.1, -0.2, -0.3]), + torch.tensor([-0.4, -0.5]), + ], + "ref_log_probs": None, + "rewards": [1.25, -0.75], + "values": None, + "response_lengths": [3, 2], + "loss_masks": [torch.ones(3), torch.ones(2)], + "total_lengths": [5, 4], + } + return component.compute_advantages_and_returns(rollout_data) + + +def test_p3o_advantages_match_grpo_shapes_and_values(): + p3o = _compute("p3o") + grpo = _compute("grpo") + + for key in ("advantages", "returns"): + p3o_values = p3o[key].unbind() + grpo_values = grpo[key].unbind() + assert [value.shape for value in p3o_values] == [torch.Size([3]), torch.Size([2])] + assert len(p3o_values) == len(grpo_values) + for p3o_value, grpo_value in zip(p3o_values, grpo_values, strict=True): + torch.testing.assert_close(p3o_value, grpo_value) + + torch.testing.assert_close(p3o["advantages"].unbind()[0], torch.full((3,), 1.25)) + torch.testing.assert_close(p3o["advantages"].unbind()[1], torch.full((2,), -0.75)) diff --git a/tests/engine/rollout/test_sglang_rollout_diagnostics.py b/tests/engine/rollout/test_sglang_rollout_diagnostics.py index 9bf06bd5b..a3728cb74 100644 --- a/tests/engine/rollout/test_sglang_rollout_diagnostics.py +++ b/tests/engine/rollout/test_sglang_rollout_diagnostics.py @@ -1,15 +1,18 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""LogprobResponse rollout-side self-topk decoding (base64 path). +"""Rollout-side log-prob decoding and multimodal pairing diagnostics. Refactored: the old module-level ``extract_sglang_topk_logprobs`` was replaced by ``LogprobResponse.self_topk("rollout", ...)`` which decodes the sglang base64 ``output_top_logprobs_*_b64`` fields into numpy ``(ids, logps)``. """ +from types import SimpleNamespace + import numpy as np import pybase64 +from relax.utils.data import processing_utils from relax.utils.opd.opd_main_worker import LogprobResponse @@ -38,3 +41,42 @@ def test_rollout_self_topk_keeps_token_id_zero_from_b64() -> None: def test_rollout_self_topk_returns_none_when_absent() -> None: assert LogprobResponse({"meta_info": {}}).self_topk("rollout", top_k=2) is None + + +def test_multimodal_token_replacement_marks_stale_behavior_logprobs() -> None: + tokenizer = SimpleNamespace( + pad_token_id=0, + image_token_id=10, + audio_token_id=11, + video_token_id=12, + ) + original = [1, 10, 2, 11, 12, 3] + + tokens, pairing_mask, counts = processing_utils._sanitize_response_tokens_for_logprobs( + tokenizer, + None, + original, + ) + + assert original == [1, 10, 2, 11, 12, 3] + assert tokens == [1, 0, 2, 0, 0, 3] + assert pairing_mask == [True, False, True, False, False, True] + assert counts == {"image": 1, "audio": 1, "video": 1} + + +def test_media_pad_replacement_marks_stale_behavior_logprob(monkeypatch) -> None: + monkeypatch.setattr( + processing_utils, + "sanitize_kimi_k25_response_tokens", + lambda processor, tokens: [tokens[0], 0, tokens[2]], + ) + + tokens, pairing_mask, counts = processing_utils._sanitize_response_tokens_for_logprobs( + SimpleNamespace(pad_token_id=0), + object(), + [1, 99, 2], + ) + + assert tokens == [1, 0, 2] + assert pairing_mask == [True, False, True] + assert counts == {"media_pad": 1} diff --git a/tests/examples/algorithms/p3o/test_configs.py b/tests/examples/algorithms/p3o/test_configs.py new file mode 100644 index 000000000..dc6f32704 --- /dev/null +++ b/tests/examples/algorithms/p3o/test_configs.py @@ -0,0 +1,719 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Static comparability tests for the P3O A100x4 launch scripts.""" + +import json +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[4] +SCRIPT_DIR = REPO_ROOT / "examples" / "algorithms" / "p3o" +FORMAL_SCRIPTS = { + "p3o_on_policy": SCRIPT_DIR / "run_p3o_on_policy_a100x4.sh", + "grpo_on_policy": SCRIPT_DIR / "run_grpo_on_policy_a100x4.sh", + "p3o_temperature_1p2": SCRIPT_DIR / "run_p3o_temperature_1p2_a100x4.sh", + "grpo_temperature_1p2": SCRIPT_DIR / "run_grpo_temperature_1p2_a100x4.sh", +} +LOW_TEMPERATURE_SCRIPTS = { + "p3o_temperature_0p6": SCRIPT_DIR / "run_p3o_temperature_0p6_a100x4.sh", + "grpo_temperature_0p6": SCRIPT_DIR / "run_grpo_temperature_0p6_a100x4.sh", +} +PERIODIC_SYNC_SCRIPTS = { + "p3o_periodic_sync_interval_3": SCRIPT_DIR / "run_p3o_periodic_sync_interval_3_a100x4.sh", + "grpo_periodic_sync_interval_3": SCRIPT_DIR / "run_grpo_periodic_sync_interval_3_a100x4.sh", +} +ALL_SCENARIO_SCRIPTS = {**FORMAL_SCRIPTS, **LOW_TEMPERATURE_SCRIPTS, **PERIODIC_SYNC_SCRIPTS} + + +def _bash_executable() -> str: + """Resolve a POSIX bash that can open the repository's own paths. + + A bare ``bash`` argv[0] is not safe to rely on: Windows resolves + executables 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. Prefer an explicit Git-for-Windows + bash, and skip rather than fail when no usable POSIX shell exists. + """ + explicit_bash_dir = os.environ.get("GIT_BASH_DIR") + candidates = [] + if explicit_bash_dir: + candidates.append(shutil.which("bash", path=explicit_bash_dir)) + if os.name == "nt": + candidates.extend( + [ + r"C:\Program Files\Git\usr\bin\bash.exe", + r"C:\Program Files\Git\bin\bash.exe", + ] + ) + else: + candidates.extend(["/bin/bash", "/usr/bin/bash", shutil.which("bash")]) + + for candidate in candidates: + if candidate and Path(candidate).is_file(): + return candidate + pytest.skip("no POSIX bash available to dry-run the launch scripts") + + +def _shell_path(path: Path, bash: str) -> str: + """Translate a Windows path for Git Bash; POSIX paths pass through.""" + if os.name != "nt": + return str(path) + del bash + normalized = path.resolve().as_posix() + return f"/{normalized[0].lower()}{normalized[2:]}" + + +def _dry_run(script: Path, *extra_args: str, env_overrides: dict[str, str] | None = None) -> list[str]: + env = os.environ.copy() + for name in ( + "P3O_ACTIVATION_RECOMPUTE", + "P3O_CLIP_HIGH", + "P3O_CLIP_LOW", + "P3O_CLEAR_RUNTIME_PROXIES", + "P3O_DETERMINISTIC_INFERENCE", + "P3O_ESS_SCOPE", + "P3O_EVAL_MAX_RESPONSE_LEN", + "P3O_EVAL_NAME", + "P3O_EVAL_N_SAMPLES", + "P3O_EVAL_TEMPERATURE", + "P3O_EVAL_TOP_P", + "P3O_INPUT_KEY", + "P3O_KL_MODE", + "P3O_LABEL_KEY", + "P3O_LOG_PROBS_CHUNK_SIZE", + "P3O_MODE", + "P3O_MODEL_CONFIG", + "P3O_MODEL_ROTARY_BASE", + "P3O_NUM_ROLLOUT", + "P3O_RM_TYPE", + "P3O_ROLLOUT_RESULT_DIR", + "P3O_ROLLOUT_SHUFFLE", + "P3O_ROLLOUT_BATCH_SIZE", + "P3O_N_SAMPLES", + ): + env.pop(name, None) + env["P3O_DRY_RUN"] = "1" + env["P3O_RAY_DASHBOARD"] = "http://example.invalid:8265" + if env_overrides is not None: + env.update(env_overrides) + bash = _bash_executable() + if os.name == "nt": + env["PATH"] = f"{Path(bash).parent}{os.pathsep}{env.get('PATH', '')}" + result = subprocess.run( + [bash, str(script), *extra_args], + cwd=REPO_ROOT, + env=env, + check=True, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + ) + return result.stdout.splitlines() + + +def _run_fake_ray( + tmp_path: Path, + script: Path, + *, + submit_exit_code: int = 0, + env_overrides: dict[str, str] | None = None, +) -> tuple[subprocess.CompletedProcess[str], Path, list[str]]: + """Run a real launcher path against a recording fake Ray executable.""" + bash = _bash_executable() + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_ray = fake_bin / "ray" + fake_ray.write_text( + "#!/bin/bash\n" + 'printf "%s\\n" "$@" >>"${FAKE_RAY_CALLS}"\n' + 'if [[ "$1 $2" == "job submit" ]]; then\n' + ' exit "${FAKE_RAY_SUBMIT_EXIT}"\n' + "fi\n" + 'if [[ "$1 $2" == "job status" ]]; then\n' + " echo TERMINAL\n" + " exit 0\n" + "fi\n" + "exit 2\n", + encoding="utf-8", + ) + fake_ray.chmod(0o755) + + model_dir = tmp_path / "model" + megatron_dir = tmp_path / "megatron" + model_dir.mkdir() + megatron_dir.mkdir() + train_data = tmp_path / "train.jsonl" + train_data.write_text("{}\n", encoding="utf-8") + output_root = tmp_path / "output" + ray_calls = tmp_path / "ray_calls.txt" + + env = os.environ.copy() + for name in ( + "P3O_ACTIVATION_RECOMPUTE", + "P3O_ALGORITHM", + "P3O_BEHAVIOR_TEMPERATURE", + "P3O_CLEAR_RUNTIME_PROXIES", + "P3O_DETERMINISTIC_INFERENCE", + "P3O_ENABLE_TEMPERATURE_OVERRIDE", + "P3O_EVAL_MAX_RESPONSE_LEN", + "P3O_EVAL_NAME", + "P3O_EVAL_N_SAMPLES", + "P3O_EVAL_TEMPERATURE", + "P3O_EVAL_TOP_P", + "P3O_LOG_PROBS_CHUNK_SIZE", + "P3O_NCCL_DEBUG", + "P3O_MODEL_CONFIG", + "P3O_MODEL_ROTARY_BASE", + "P3O_RM_TYPE", + "P3O_ROLLOUT_RESULT_DIR", + "P3O_ROLLOUT_SHUFFLE", + "P3O_TORCH_DISTRIBUTED_DEBUG", + "P3O_UPDATE_WEIGHTS_INTERVAL", + ): + env.pop(name, None) + env.update( + { + "FAKE_RAY_CALLS": str(ray_calls), + "FAKE_RAY_SUBMIT_EXIT": str(submit_exit_code), + "P3O_DRY_RUN": "0", + "P3O_MEGATRON_DIR": str(megatron_dir), + "P3O_MODE": "smoke", + "P3O_MODEL_DIR": str(model_dir), + "P3O_OUTPUT_ROOT": str(output_root), + "P3O_RAY_DASHBOARD": "http://example.invalid:8265", + "P3O_RUN_ID": "integration", + "P3O_TRAIN_DATA": str(train_data), + } + ) + if env_overrides is not None: + env.update(env_overrides) + if os.name == "nt": + env["PATH"] = f"{Path(bash).parent}{os.pathsep}{env.get('PATH', '')}" + for name in ( + "FAKE_RAY_CALLS", + "P3O_EVAL_DATA", + "P3O_MEGATRON_DIR", + "P3O_MODEL_DIR", + "P3O_OUTPUT_ROOT", + "P3O_TRAIN_DATA", + ): + if name in env: + env[name] = _shell_path(Path(env[name]), bash) + + result = subprocess.run( + [ + bash, + "-c", + 'export PATH="$1:$PATH"; exec "$2"', + "p3o-runner", + _shell_path(fake_bin, bash), + _shell_path(script, bash), + ], + cwd=REPO_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + ) + config_name = script.stem.removeprefix("run_").removesuffix("_a100x4") + run_dir = output_root / config_name / "seed_42" / "integration" + calls = ray_calls.read_text(encoding="utf-8").splitlines() if ray_calls.exists() else [] + return result, run_dir, calls + + +def _option_value(args: list[str], option: str) -> str: + return args[args.index(option) + 1] + + +def _comparable_args(args: list[str]) -> list[str]: + ignored_with_value = { + "--advantage-estimator", + "--eps-clip", + "--eps-clip-high", + "--p3o-ess-scope", + "--p3o-kl-mode", + "--clip-low", + "--clip-high", + "--tb-experiment-name", + } + normalized = [] + index = 0 + while index < len(args): + if args[index] in ignored_with_value: + index += 2 + else: + normalized.append(args[index]) + index += 1 + return normalized + + +def test_p3o_configs_freeze_required_formal_values(): + for args in map(_dry_run, FORMAL_SCRIPTS.values()): + assert _option_value(args, "--input-key") == "problem" + assert _option_value(args, "--label-key") == "answer" + assert _option_value(args, "--rm-type") == "deepscaler" + assert _option_value(args, "--num-rollout") == "30" + assert _option_value(args, "--rollout-batch-size") == "4" + assert _option_value(args, "--n-samples-per-prompt") == "16" + assert _option_value(args, "--global-batch-size") == "64" + assert _option_value(args, "--micro-batch-size") == "1" + assert _option_value(args, "--rollout-max-response-len") == "4096" + assert _option_value(args, "--rollout-temperature") == "1.0" + assert _option_value(args, "--rollout-top-p") == "1.0" + assert _option_value(args, "--lr") == "1e-5" + assert _option_value(args, "--adam-beta2") == "0.95" + assert _option_value(args, "--weight-decay") == "0.01" + assert "--calculate-per-token-loss" in args + assert "--use-rollout-logprobs" in args + assert "--rollout-shuffle" in args + assert "--colocate" in args + assert "--fully-async" not in args + assert "--use-tis" not in args + assert "--use-kl-loss" not in args + assert "--eval-size" not in args + assert _option_value(args, "--num-layers") == "36" + assert _option_value(args, "--hidden-size") == "2560" + assert _option_value(args, "--rotary-base") == "5000000" + assert ( + int(_option_value(args, "--num-rollout")) + * int(_option_value(args, "--rollout-batch-size")) + * int(_option_value(args, "--n-samples-per-prompt")) + == 1920 + ) + assert int(_option_value(args, "--rollout-batch-size")) % 4 == 0 + assert int(_option_value(args, "--global-batch-size")) == ( + int(_option_value(args, "--rollout-batch-size")) * int(_option_value(args, "--n-samples-per-prompt")) + ) + + +def test_p3o_configs_use_active_algorithm_settings_only_for_p3o(): + p3o_args = _dry_run(FORMAL_SCRIPTS["p3o_on_policy"]) + grpo_args = _dry_run(FORMAL_SCRIPTS["grpo_on_policy"]) + + assert _option_value(p3o_args, "--p3o-ess-scope") == "micro-batch" + assert _option_value(p3o_args, "--p3o-kl-mode") == "proxy_safe" + assert _option_value(p3o_args, "--clip-low") == "0.2" + assert _option_value(p3o_args, "--clip-high") == "0.2" + for option in ("--p3o-ess-scope", "--p3o-kl-mode", "--clip-low", "--clip-high"): + assert option not in grpo_args + + +def test_p3o_configs_are_comparable_within_each_scenario(): + resolved = {name: _dry_run(script) for name, script in FORMAL_SCRIPTS.items()} + assert _comparable_args(resolved["p3o_on_policy"]) == _comparable_args(resolved["grpo_on_policy"]) + assert _comparable_args(resolved["p3o_temperature_1p2"]) == _comparable_args(resolved["grpo_temperature_1p2"]) + + assert "--custom-generate-function-path" not in resolved["p3o_on_policy"] + assert "--custom-generate-function-path" not in resolved["grpo_on_policy"] + for name in ("p3o_temperature_1p2", "grpo_temperature_1p2"): + assert _option_value(resolved[name], "--custom-generate-function-path") == ( + "examples.algorithms.p3o.rollout.generate" + ) + + for name in ("grpo_on_policy", "grpo_temperature_1p2"): + assert _option_value(resolved[name], "--eps-clip") == "0.4" + assert _option_value(resolved[name], "--eps-clip-high") == "0.4" + for name in ("p3o_on_policy", "p3o_temperature_1p2"): + assert "--eps-clip" not in resolved[name] + assert "--eps-clip-high" not in resolved[name] + + +def test_p3o_low_temperature_configs_are_matched_and_named_from_temperature(): + resolved = {name: _dry_run(script) for name, script in LOW_TEMPERATURE_SCRIPTS.items()} + + p3o_args = resolved["p3o_temperature_0p6"] + grpo_args = resolved["grpo_temperature_0p6"] + assert _option_value(p3o_args, "--tb-experiment-name") == "p3o_temperature_0p6-seed-42" + assert _option_value(grpo_args, "--tb-experiment-name") == "grpo_temperature_0p6-seed-42" + assert _option_value(p3o_args, "--update-weights-interval") == "1" + assert _option_value(grpo_args, "--update-weights-interval") == "1" + assert _option_value(p3o_args, "--custom-generate-function-path") == ("examples.algorithms.p3o.rollout.generate") + assert _comparable_args(p3o_args) == _comparable_args(grpo_args) + + smoke_args = _dry_run(SCRIPT_DIR / "run_p3o_smoke.sh", "p3o_temperature_0p6") + assert _option_value(smoke_args, "--tb-experiment-name") == "p3o_temperature_0p6-seed-42" + + +def test_p3o_periodic_sync_configs_are_matched_and_parameterized(): + resolved = {name: _dry_run(script) for name, script in PERIODIC_SYNC_SCRIPTS.items()} + + p3o_args = resolved["p3o_periodic_sync_interval_3"] + grpo_args = resolved["grpo_periodic_sync_interval_3"] + assert _option_value(p3o_args, "--max-staleness") == "0" + assert _option_value(grpo_args, "--max-staleness") == "0" + assert _option_value(p3o_args, "--update-weights-interval") == "3" + assert _option_value(grpo_args, "--update-weights-interval") == "3" + assert _option_value(p3o_args, "--tb-experiment-name") == "p3o_periodic_sync_interval_3-seed-42" + assert _option_value(grpo_args, "--tb-experiment-name") == "grpo_periodic_sync_interval_3-seed-42" + assert _comparable_args(p3o_args) == _comparable_args(grpo_args) + + overridden = _dry_run( + PERIODIC_SYNC_SCRIPTS["p3o_periodic_sync_interval_3"], + env_overrides={"P3O_UPDATE_WEIGHTS_INTERVAL": "5"}, + ) + assert _option_value(overridden, "--max-staleness") == "0" + assert _option_value(overridden, "--update-weights-interval") == "5" + assert _option_value(overridden, "--tb-experiment-name") == "p3o_periodic_sync_interval_5-seed-42" + + +def test_p3o_smoke_uses_one_small_optimizer_step(): + args = _dry_run(SCRIPT_DIR / "run_p3o_smoke.sh", "p3o_temperature_1p2") + + assert _option_value(args, "--num-rollout") == "1" + assert _option_value(args, "--rollout-batch-size") == "4" + assert _option_value(args, "--n-samples-per-prompt") == "4" + assert _option_value(args, "--global-batch-size") == "16" + assert _option_value(args, "--micro-batch-size") == "1" + assert _option_value(args, "--rollout-max-response-len") == "128" + assert _option_value(args, "--input-key") == "question" + assert _option_value(args, "--rm-type") == "mopd" + assert _option_value(args, "--num-layers") == "28" + assert _option_value(args, "--hidden-size") == "1024" + assert _option_value(args, "--rotary-base") == "1000000" + assert _option_value(args, "--p3o-ess-scope") == "micro-batch" + assert _option_value(args, "--p3o-kl-mode") == "proxy_safe" + assert _option_value(args, "--rollout-result-dir") == "/dummy/output/rollout_results" + assert "--eval-prompt-data" not in args + + +def test_p3o_dataset_keys_can_be_overridden(): + args = _dry_run( + SCRIPT_DIR / "run_p3o_smoke.sh", + "p3o_on_policy", + env_overrides={"P3O_INPUT_KEY": "problem", "P3O_LABEL_KEY": "solution"}, + ) + + assert _option_value(args, "--input-key") == "problem" + assert _option_value(args, "--label-key") == "solution" + + +def test_p3o_reward_type_can_be_overridden_for_deepscaler_smoke(): + args = _dry_run( + SCRIPT_DIR / "run_p3o_smoke.sh", + "p3o_on_policy", + env_overrides={"P3O_RM_TYPE": "deepscaler"}, + ) + + assert _option_value(args, "--rm-type") == "deepscaler" + + +def test_p3o_rollout_result_dir_can_be_overridden(): + args = _dry_run( + SCRIPT_DIR / "run_p3o_smoke.sh", + "p3o_on_policy", + env_overrides={"P3O_ROLLOUT_RESULT_DIR": "/evidence/raw_rollouts"}, + ) + + assert _option_value(args, "--rollout-result-dir") == "/evidence/raw_rollouts" + + +def test_p3o_rollout_shuffle_can_be_disabled_for_a_fixed_prompt_schedule(): + args = _dry_run( + SCRIPT_DIR / "run_p3o_smoke.sh", + "p3o_on_policy", + env_overrides={"P3O_ROLLOUT_SHUFFLE": "0"}, + ) + + assert "--rollout-shuffle" not in args + + +def test_p3o_deterministic_inference_can_be_enabled_for_paired_sampling(): + args = _dry_run( + SCRIPT_DIR / "run_p3o_smoke.sh", + "p3o_on_policy", + env_overrides={"P3O_DETERMINISTIC_INFERENCE": "1"}, + ) + + assert args.count("--sglang-enable-deterministic-inference") == 1 + + +def test_p3o_smoke_can_select_pipeline_parallel_size_two(): + args = _dry_run( + SCRIPT_DIR / "run_p3o_smoke.sh", + "p3o_on_policy", + env_overrides={"P3O_PIPELINE_MODEL_PARALLEL_SIZE": "2", "P3O_NUM_ROLLOUT": "3"}, + ) + + assert _option_value(args, "--pipeline-model-parallel-size") == "2" + assert _option_value(args, "--num-rollout") == "3" + assert _option_value(args, "--tb-experiment-name") == "p3o_on_policy_pp2-seed-42" + + +def test_p3o_runner_requires_explicit_ray_dashboard(): + env = os.environ.copy() + env.pop("P3O_RAY_DASHBOARD", None) + env["P3O_DRY_RUN"] = "1" + bash = _bash_executable() + if os.name == "nt": + env["PATH"] = f"{Path(bash).parent}{os.pathsep}{env.get('PATH', '')}" + result = subprocess.run( + [bash, str(FORMAL_SCRIPTS["p3o_on_policy"])], + cwd=REPO_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + ) + + assert result.returncode != 0 + assert "P3O_RAY_DASHBOARD must be set" in result.stderr + + +@pytest.mark.parametrize( + ("scenario", "expected_algorithm", "expected_interval", "expected_temperature"), + [ + ("p3o_on_policy", "p3o", "1", None), + ("grpo_on_policy", "grpo", "1", None), + ("p3o_periodic_sync_interval_3", "p3o", "3", None), + ("grpo_periodic_sync_interval_3", "grpo", "3", None), + ("p3o_temperature_0p6", "p3o", "1", "0.6"), + ("grpo_temperature_0p6", "grpo", "1", "0.6"), + ("p3o_temperature_1p2", "p3o", "1", "1.2"), + ("grpo_temperature_1p2", "grpo", "1", "1.2"), + ], +) +def test_p3o_runner_executes_all_scenarios_with_fake_ray( + tmp_path, + scenario, + expected_algorithm, + expected_interval, + expected_temperature, +): + result, run_dir, ray_calls = _run_fake_ray(tmp_path, ALL_SCENARIO_SCRIPTS[scenario]) + + assert result.returncode == 0, result.stderr + resolved_args = (run_dir / "resolved_args.txt").read_text(encoding="utf-8").splitlines() + runtime_env = json.loads(_option_value(ray_calls, "--runtime-env-json"))["env_vars"] + identity = dict( + line.split("=", 1) + for line in (run_dir / "run_identity.env").read_text(encoding="utf-8").splitlines() + if "=" in line + ) + assert _option_value(resolved_args, "--advantage-estimator") == expected_algorithm + assert _option_value(resolved_args, "--update-weights-interval") == expected_interval + assert _option_value(ray_calls, "--submission-id") == f"{scenario}-seed-42-integration" + assert runtime_env["NCCL_DEBUG"] == "WARN" + assert runtime_env["TORCH_DISTRIBUTED_DEBUG"] == "OFF" + assert runtime_env["RAY_OVERRIDE_JOB_RUNTIME_ENV"] == "1" + for proxy_name in ( + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "http_proxy", + "https_proxy", + "all_proxy", + "NO_PROXY", + "no_proxy", + ): + assert proxy_name not in runtime_env + assert {"GIT_COMMIT", "GIT_BRANCH", "GIT_DIRTY", "started_utc", "ended_utc"} <= identity.keys() + assert identity["model_config"].endswith("qwen3-0.6B.sh") + assert identity["model_rotary_base"] == _option_value(resolved_args, "--rotary-base") == "1000000" + assert identity["p3o_ess_scope"] == "micro-batch" + assert identity["p3o_kl_mode"] == "proxy_safe" + assert identity["clip_low"] == identity["clip_high"] == "0.2" + assert identity["input_key"] == "question" + assert identity["label_key"] == "answer" + assert identity["rm_type"] == "mopd" + assert identity["rollout_shuffle"] == "1" + assert identity["clear_runtime_proxies"] == "0" + assert identity["rollout_result_dir"].endswith(f"/{scenario}/seed_42/integration/rollout_results") + assert _option_value(resolved_args, "--rollout-result-dir").endswith( + f"/{scenario}/seed_42/integration/rollout_results" + ) + assert identity["config"] == scenario + assert identity["ray_job_id"] == f"{scenario}-seed-42-integration" + if expected_temperature is None: + assert "--custom-generate-function-path" not in resolved_args + assert "P3O_BEHAVIOR_TEMPERATURE" not in runtime_env + else: + assert _option_value(resolved_args, "--custom-generate-function-path") == ( + "examples.algorithms.p3o.rollout.generate" + ) + assert runtime_env["P3O_BEHAVIOR_TEMPERATURE"] == expected_temperature + assert (run_dir / "exit_code.txt").read_text(encoding="utf-8").strip() == "0" + assert (run_dir / "job_status.txt").read_text(encoding="utf-8").strip() == "TERMINAL" + + +def test_p3o_runner_can_clear_runtime_proxies_explicitly(tmp_path): + result, run_dir, ray_calls = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={"P3O_CLEAR_RUNTIME_PROXIES": "1"}, + ) + + assert result.returncode == 0, result.stderr + runtime_env = json.loads(_option_value(ray_calls, "--runtime-env-json"))["env_vars"] + identity = (run_dir / "run_identity.env").read_text(encoding="utf-8") + for proxy_name in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): + assert runtime_env[proxy_name] == "" + assert runtime_env["NO_PROXY"] == runtime_env["no_proxy"] == "*" + assert "clear_runtime_proxies=1" in identity + + +def test_p3o_runner_preserves_debug_overrides(tmp_path): + result, _, ray_calls = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={"P3O_NCCL_DEBUG": "INFO", "P3O_TORCH_DISTRIBUTED_DEBUG": "DETAIL"}, + ) + + assert result.returncode == 0, result.stderr + runtime_env = json.loads(_option_value(ray_calls, "--runtime-env-json"))["env_vars"] + assert runtime_env["NCCL_DEBUG"] == "INFO" + assert runtime_env["TORCH_DISTRIBUTED_DEBUG"] == "DETAIL" + + +def test_p3o_runner_preserves_failed_ray_exit_code(tmp_path): + result, run_dir, _ = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + submit_exit_code=17, + ) + + assert result.returncode == 17 + assert (run_dir / "exit_code.txt").read_text(encoding="utf-8").strip() == "17" + assert (run_dir / "job_status.txt").read_text(encoding="utf-8").strip() == "TERMINAL" + + +def test_p3o_smoke_runner_does_not_require_eval_data(tmp_path): + result, _, _ = _run_fake_ray(tmp_path, FORMAL_SCRIPTS["p3o_on_policy"]) + + assert result.returncode == 0, result.stderr + + +def test_p3o_formal_runner_requires_eval_data(tmp_path): + result, _, ray_calls = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={"P3O_MODE": "formal"}, + ) + + assert result.returncode == 1 + assert "P3O_EVAL_DATA must be set in formal mode" in result.stderr + assert ray_calls == [] + + +def test_p3o_formal_runner_accepts_existing_eval_data(tmp_path): + eval_data = tmp_path / "eval.jsonl" + eval_data.write_text("{}\n", encoding="utf-8") + + result, run_dir, _ = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={"P3O_MODE": "formal", "P3O_EVAL_DATA": str(eval_data)}, + ) + + assert result.returncode == 0, result.stderr + resolved_args = (run_dir / "resolved_args.txt").read_text(encoding="utf-8").splitlines() + assert _option_value(resolved_args, "--eval-prompt-data") == "deepscaler" + assert str(eval_data.name) in resolved_args[resolved_args.index("--eval-prompt-data") + 2] + assert _option_value(resolved_args, "--n-samples-per-eval-prompt") == "16" + assert _option_value(resolved_args, "--eval-max-response-len") == "4096" + assert _option_value(resolved_args, "--eval-temperature") == "1.0" + assert _option_value(resolved_args, "--eval-top-p") == "0.95" + + +def test_p3o_formal_runner_records_resource_adjusted_eval_contract(tmp_path): + eval_data = tmp_path / "eval.jsonl" + eval_data.write_text("{}\n", encoding="utf-8") + + result, run_dir, _ = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={ + "P3O_MODE": "formal", + "P3O_EVAL_DATA": str(eval_data), + "P3O_EVAL_NAME": "local-deepscaler", + "P3O_EVAL_N_SAMPLES": "1", + "P3O_EVAL_MAX_RESPONSE_LEN": "2048", + "P3O_EVAL_TEMPERATURE": "0.8", + "P3O_EVAL_TOP_P": "0.9", + }, + ) + + assert result.returncode == 0, result.stderr + resolved_args = (run_dir / "resolved_args.txt").read_text(encoding="utf-8").splitlines() + identity = (run_dir / "run_identity.env").read_text(encoding="utf-8") + assert _option_value(resolved_args, "--eval-prompt-data") == "local-deepscaler" + assert _option_value(resolved_args, "--n-samples-per-eval-prompt") == "1" + assert _option_value(resolved_args, "--eval-max-response-len") == "2048" + assert _option_value(resolved_args, "--eval-temperature") == "0.8" + assert _option_value(resolved_args, "--eval-top-p") == "0.9" + assert "eval_name=local-deepscaler" in identity + assert "eval_n_samples=1" in identity + assert "eval_max_response_len=2048" in identity + + +def test_p3o_runner_records_resource_adjusted_activation_contract(tmp_path): + result, run_dir, _ = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={ + "P3O_ACTIVATION_RECOMPUTE": "1", + "P3O_LOG_PROBS_CHUNK_SIZE": "128", + }, + ) + + assert result.returncode == 0, result.stderr + resolved_args = (run_dir / "resolved_args.txt").read_text(encoding="utf-8").splitlines() + identity = (run_dir / "run_identity.env").read_text(encoding="utf-8") + assert _option_value(resolved_args, "--recompute-granularity") == "full" + assert _option_value(resolved_args, "--recompute-method") == "uniform" + assert _option_value(resolved_args, "--recompute-num-layers") == "1" + assert _option_value(resolved_args, "--log-probs-chunk-size") == "128" + assert "activation_recompute=1" in identity + assert "log_probs_chunk_size=128" in identity + + +def test_p3o_runner_validates_megatron_directory_before_ray(tmp_path): + missing_megatron = tmp_path / "missing-megatron" + result, _, ray_calls = _run_fake_ray( + tmp_path, + FORMAL_SCRIPTS["p3o_on_policy"], + env_overrides={"P3O_MEGATRON_DIR": str(missing_megatron)}, + ) + + assert result.returncode == 2 + assert missing_megatron.name in result.stderr + assert ray_calls == [] + + +@pytest.mark.parametrize("raw_value", ["", "0", "0.0", "-1", "NaN", "Inf", "warm"]) +def test_p3o_shell_rejects_invalid_behavior_temperature(raw_value): + env = os.environ.copy() + env.update( + { + "P3O_ALGORITHM": "p3o", + "P3O_BEHAVIOR_TEMPERATURE": raw_value, + "P3O_DRY_RUN": "1", + "P3O_ENABLE_TEMPERATURE_OVERRIDE": "1", + "P3O_RAY_DASHBOARD": "http://example.invalid:8265", + } + ) + bash = _bash_executable() + if os.name == "nt": + env["PATH"] = f"{Path(bash).parent}{os.pathsep}{env.get('PATH', '')}" + result = subprocess.run( + [bash, "-c", 'source "$1"; P3O_run', "p3o-test", str(SCRIPT_DIR / "common_a100x4.sh")], + cwd=REPO_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + ) + + assert result.returncode != 0 + assert "P3O_BEHAVIOR_TEMPERATURE" in result.stderr diff --git a/tests/examples/algorithms/p3o/test_rollout.py b/tests/examples/algorithms/p3o/test_rollout.py new file mode 100644 index 000000000..02e1abba6 --- /dev/null +++ b/tests/examples/algorithms/p3o/test_rollout.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for the P3O behavior-only temperature wrapper.""" + +from types import SimpleNamespace + +import pytest + +from examples.algorithms.p3o import rollout + + +def test_behavior_sampling_params_overrides_only_temperature(monkeypatch): + monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", "0.6") + original = {"temperature": 1.0, "top_p": 0.9, "max_new_tokens": 64} + + updated = rollout.behavior_sampling_params(original, evaluation=False) + + assert updated == {"temperature": 0.6, "top_p": 0.9, "max_new_tokens": 64} + assert original == {"temperature": 1.0, "top_p": 0.9, "max_new_tokens": 64} + + +@pytest.mark.parametrize(("raw_value", "expected"), [("0.6", 0.6), ("1.2", 1.2), ("2.0", 2.0)]) +def test_behavior_sampling_params_accepts_runtime_temperature(monkeypatch, raw_value, expected): + monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", raw_value) + + updated = rollout.behavior_sampling_params({"temperature": 1.0}, evaluation=False) + + assert updated["temperature"] == expected + + +def test_behavior_sampling_params_requires_runtime_temperature(monkeypatch): + monkeypatch.delenv("P3O_BEHAVIOR_TEMPERATURE", raising=False) + + with pytest.raises(ValueError, match="must be set"): + rollout.behavior_sampling_params({"temperature": 1.0}, evaluation=False) + + +@pytest.mark.parametrize("raw_value", ["0", "0.0", "-1", "nan", "inf", "-inf"]) +def test_behavior_sampling_params_rejects_non_positive_or_nonfinite_temperature(monkeypatch, raw_value): + monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", raw_value) + + with pytest.raises(ValueError, match="finite and greater than zero"): + rollout.behavior_sampling_params({"temperature": 1.0}, evaluation=False) + + +def test_behavior_sampling_params_rejects_nonnumeric_temperature(monkeypatch): + monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", "warm") + + with pytest.raises(ValueError, match="must be numeric"): + rollout.behavior_sampling_params({"temperature": 1.0}, evaluation=False) + + +def test_behavior_sampling_params_preserves_evaluation_without_temperature_env(monkeypatch): + monkeypatch.delenv("P3O_BEHAVIOR_TEMPERATURE", raising=False) + original = {"temperature": 0.0, "top_p": 0.7, "max_new_tokens": 128} + + updated = rollout.behavior_sampling_params(original, evaluation=True) + + assert updated == original + assert updated is not original + + +async def test_generate_delegates_with_isolated_behavior_params(monkeypatch): + monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", "1.2") + captured = {} + expected = object() + + async def fake_generate(args, sample, sampling_params, evaluation=False): + captured.update( + args=args, + sample=sample, + sampling_params=sampling_params, + evaluation=evaluation, + ) + return expected + + monkeypatch.setattr(rollout, "_sglang_generate", fake_generate) + args = SimpleNamespace() + sample = object() + original = {"temperature": 1.0, "top_p": 0.95, "max_new_tokens": 32} + + result = await rollout.generate(args, sample, original, evaluation=False) + + assert result is expected + assert captured == { + "args": args, + "sample": sample, + "sampling_params": {"temperature": 1.2, "top_p": 0.95, "max_new_tokens": 32}, + "evaluation": False, + } + assert original == {"temperature": 1.0, "top_p": 0.95, "max_new_tokens": 32} diff --git a/tests/utils/test_arguments_opd_teacher_colocate.py b/tests/utils/test_arguments_opd_teacher_colocate.py index b87721536..ddf352007 100644 --- a/tests/utils/test_arguments_opd_teacher_colocate.py +++ b/tests/utils/test_arguments_opd_teacher_colocate.py @@ -194,6 +194,14 @@ def test_opd_sampled_token_loss_is_accepted(arguments_module): arguments_module.slime_validate_args(args) +def test_p3o_with_opd_is_rejected_before_training(arguments_module): + args = _opd_args() + args.advantage_estimator = "p3o" + + with pytest.raises(ValueError, match="P3O and OPD are mutually exclusive"): + arguments_module.slime_validate_args(args) + + def test_managed_opd_teacher_colocate_preserves_rollout_resource_split(arguments_module): args = _opd_args() args.colocate = True diff --git a/tests/utils/test_multimodal_rollout_stats.py b/tests/utils/test_multimodal_rollout_stats.py index 4bb25afe8..610378794 100644 --- a/tests/utils/test_multimodal_rollout_stats.py +++ b/tests/utils/test_multimodal_rollout_stats.py @@ -41,6 +41,8 @@ def test_rollout_summary_record_includes_token_and_agent_stats(): response="world", tokens=list(range(12)), response_length=5, + rollout_log_probs=[-0.5, -0.4, -0.3, -0.2, -0.1], + rollout_log_probs_mask=[True, True, False, True, True], reward=1.0, multimodal_inputs={"images": ["image.png"]}, multimodal_train_inputs={"image_grid_thw": [[1, 8, 8]]}, @@ -51,6 +53,9 @@ def test_rollout_summary_record_includes_token_and_agent_stats(): assert record["prompt_token_count"] == 7 assert record["response_token_count"] == 5 + assert record["response_token_ids"] == [7, 8, 9, 10, 11] + assert record["response_rollout_log_probs"] == [-0.5, -0.4, -0.3, -0.2, -0.1] + assert record["response_rollout_log_probs_mask"] == [True, True, False, True, True] assert record["total_token_count"] == 12 assert record["prompt_length"] == 7 assert record["image_count"] == 1 diff --git a/tests/utils/test_p3o_arguments.py b/tests/utils/test_p3o_arguments.py new file mode 100644 index 000000000..230836a29 --- /dev/null +++ b/tests/utils/test_p3o_arguments.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for the P3O configuration gates in ``arguments.py``. + +Every rejection below guards a config that still *trains* -- it just silently +optimizes something other than the P3O objective (uncorrected ratio, per- +micro-batch denominator, double correction) or breaks the pre-pass replay +(FP8 amax history, dropout). A plausible loss curve is the failure mode, so +these are hard errors rather than warnings and are worth pinning. + +``relax.utils.arguments`` pulls in the Megatron/Ray import chain, which is not +available in the unit-test environment, so the validator is extracted from the +module source by AST rather than imported. +""" + +import ast +import types +from argparse import Namespace +from pathlib import Path + +import pytest + + +ARGUMENTS_PATH = Path(__file__).resolve().parents[2] / "relax" / "utils" / "arguments.py" + + +def _load_validator(): + """Extract ``_validate_p3o_args`` without importing arguments.py.""" + import argparse + + tree = ast.parse(ARGUMENTS_PATH.read_text(encoding="utf-8")) + func = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_validate_p3o_args") + module = types.ModuleType("_p3o_args") + module.argparse = argparse # Inject argparse for type annotation + exec(compile(ast.Module(body=[func], type_ignores=[]), str(ARGUMENTS_PATH), "exec"), module.__dict__) + return module._validate_p3o_args + + +validate_p3o_args = _load_validator() + + +def _p3o_kl_mode_choices() -> list[str]: + """Read the CLI choices without importing Relax's Megatron dependency + chain.""" + tree = ast.parse(ARGUMENTS_PATH.read_text(encoding="utf-8")) + call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "add_argument" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "--p3o-kl-mode" + ) + choices = next(keyword.value for keyword in call.keywords if keyword.arg == "choices") + return ast.literal_eval(choices) + + +def _p3o_args(**overrides) -> Namespace: + """A minimal P3O-valid config, with individual fields overridable.""" + config = dict( + advantage_estimator="p3o", + p3o_ess_scope="micro-batch", + p3o_kl_mode="proxy", + clip_low=0.2, + clip_high=0.2, + use_rollout_logprobs=True, + calculate_per_token_loss=True, + use_tis=False, + true_on_policy_mode=False, + use_critic=False, + fp8=None, + attention_dropout=0.0, + hidden_dropout=0.0, + lora_rank=0, + lora_dropout=0.0, + fully_async=False, + get_mismatch_metrics=False, + use_opsm=False, + custom_pg_loss_reducer_function_path=None, + enable_mtp_training=False, + use_routing_replay=False, + use_rollout_routing_replay=False, + overlap_moe_expert_parallel_comm=False, + ) + config.update(overrides) + return Namespace(**config) + + +def test_p3o_arguments_accepts_a_valid_configuration(): + validate_p3o_args(_p3o_args()) + + +def test_p3o_arguments_accepts_true_on_policy_scheduling(): + validate_p3o_args(_p3o_args(true_on_policy_mode=True)) + + +def test_p3o_arguments_rejects_non_production_kl_modes(): + with pytest.raises(ValueError, match="proxy or proxy_safe"): + validate_p3o_args(_p3o_args(p3o_kl_mode="exact")) + + +def test_p3o_arguments_expose_only_production_kl_modes(): + assert _p3o_kl_mode_choices() == ["proxy", "proxy_safe"] + + +@pytest.mark.parametrize( + "overrides", + [ + dict(fp8="hybrid"), + dict(attention_dropout=0.1), + dict(hidden_dropout=0.1), + dict(lora_rank=8, lora_dropout=0.1), + dict(fully_async=True), + ], +) +def test_p3o_arguments_micro_batch_scope_accepts_replay_sensitive_features(overrides): + validate_p3o_args(_p3o_args(**overrides)) + + +@pytest.mark.parametrize( + "overrides", + [ + dict(fp8="hybrid"), + dict(attention_dropout=0.1), + dict(hidden_dropout=0.1), + dict(lora_rank=8, lora_dropout=0.1), + dict(fully_async=True), + ], +) +def test_p3o_arguments_step_scope_rejects_replay_sensitive_features(overrides): + with pytest.raises(ValueError): + validate_p3o_args(_p3o_args(p3o_ess_scope="step", **overrides)) + + +def test_p3o_arguments_step_scope_accepts_inactive_lora_dropout(): + validate_p3o_args(_p3o_args(p3o_ess_scope="step", lora_rank=0, lora_dropout=0.1)) + + +@pytest.mark.parametrize( + ("reason", "overrides"), + [ + ("behavior policy would be undefined", dict(use_rollout_logprobs=False)), + ("per-sample-mean reintroduces a micro-batch denominator", dict(calculate_per_token_loss=False)), + ("TIS double-corrects the same mismatch", dict(use_tis=True)), + ("P3O is critic-free", dict(use_critic=True)), + ("mismatch metrics add an unverified extra forward", dict(get_mismatch_metrics=True)), + ("OPSM changes the policy-gradient mask", dict(use_opsm=True)), + ( + "custom reducer may change token-sum normalization", + dict(custom_pg_loss_reducer_function_path="pkg.reducer"), + ), + ("MTP changes forward state between replay passes", dict(enable_mtp_training=True)), + ("training routing replay changes the replayed forward", dict(use_routing_replay=True)), + ("rollout routing replay changes the replayed forward", dict(use_rollout_routing_replay=True)), + ("combined 1F1B bypasses the standard forward", dict(overlap_moe_expert_parallel_comm=True)), + ], +) +def test_p3o_arguments_rejects_configs_that_change_the_objective(reason, overrides): + with pytest.raises((AssertionError, ValueError)): + validate_p3o_args(_p3o_args(**overrides)) + + +@pytest.mark.parametrize( + "overrides", + [ + dict(p3o_ess_scope="window"), + dict(p3o_kl_mode="unsafe"), + dict(clip_low=-0.1), + dict(clip_high=-0.1), + ], +) +def test_p3o_arguments_rejects_invalid_active_plan_values(overrides): + with pytest.raises(ValueError): + validate_p3o_args(_p3o_args(**overrides)) + + +def test_p3o_arguments_validate_after_effective_value_overrides(): + tree = ast.parse(ARGUMENTS_PATH.read_text(encoding="utf-8")) + validator = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "slime_validate_args" + ) + calls = [ + node + for node in ast.walk(validator) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_validate_p3o_args" + ] + assert len(calls) == 1 + + custom_config_if = next( + node + for node in ast.walk(validator) + if isinstance(node, ast.If) + and any( + isinstance(child, ast.Attribute) and child.attr == "custom_config_path" for child in ast.walk(node.test) + ) + ) + rollout_routing_if = next( + node + for node in ast.walk(validator) + if isinstance(node, ast.If) + and any( + isinstance(child, ast.Attribute) and child.attr == "use_rollout_routing_replay" + for child in ast.walk(node.test) + ) + ) + assert calls[0].lineno > custom_config_if.end_lineno + assert calls[0].lineno > rollout_routing_if.end_lineno diff --git a/tests/utils/test_p3o_registry.py b/tests/utils/test_p3o_registry.py new file mode 100644 index 000000000..fb9deb359 --- /dev/null +++ b/tests/utils/test_p3o_registry.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Registration and rollout reward-path tests for P3O.""" + +import argparse +import math +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +# `relax.core.registry` eagerly imports `relax.components.advantages`, which imports +# `megatron.core` at module level. CI installs no megatron, so the import runs under +# the shared stub; the registry mapping and reward path under test are pure Python. +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "backends" / "megatron")) + +from _megatron_stub import stubbed_megatron_modules # noqa: E402 + + +with stubbed_megatron_modules( + ("megatron", "ray", "tensordict", "transfer_queue", "sglang", "sglang_router", "pybase64") +): + from relax.core.registry import ALGOS # noqa: E402 + from relax.utils.arguments import get_slime_extra_args_provider # noqa: E402 + from relax.utils.types import Sample # noqa: E402 + from relax.utils.utils import post_process_rewards # noqa: E402 + + +def test_p3o_registry_parser_accepts_estimator(): + parser = get_slime_extra_args_provider()(argparse.ArgumentParser()) + action = next(action for action in parser._actions if action.dest == "advantage_estimator") + + assert "p3o" in action.choices + parsed, unknown = parser.parse_known_args(["--advantage-estimator", "p3o"]) + assert parsed.advantage_estimator == "p3o" + assert unknown == [] + + +def test_p3o_registry_parser_exposes_active_plan_defaults_and_modes(): + parser = get_slime_extra_args_provider()(argparse.ArgumentParser()) + + defaults, unknown = parser.parse_known_args([]) + configured, configured_unknown = parser.parse_known_args( + [ + "--p3o-ess-scope", + "step", + "--p3o-kl-mode", + "proxy_safe", + "--clip-low", + "0.1", + "--clip-high", + "0.3", + ] + ) + + assert unknown == configured_unknown == [] + assert defaults.p3o_ess_scope == "micro-batch" + assert defaults.p3o_kl_mode == "proxy" + assert defaults.clip_low == defaults.clip_high == 0.2 + assert configured.p3o_ess_scope == "step" + assert configured.p3o_kl_mode == "proxy_safe" + assert configured.clip_low == 0.1 + assert configured.clip_high == 0.3 + + +def test_p3o_registry_uses_grpo_service_roles(): + assert "p3o" in ALGOS + assert ALGOS["p3o"].keys() == ALGOS["grpo"].keys() + for role in ALGOS["grpo"]: + assert ALGOS["p3o"][role] is ALGOS["grpo"][role] + + +def _normalized_rewards( + estimator: str, + *, + rewards: tuple[float, ...] = (1.0, 3.0, 2.0, 6.0), + n_samples_per_prompt: int = 2, + grpo_std_normalization: bool = True, +): + args = SimpleNamespace( + custom_reward_post_process_path=None, + agentic_custom_advantage_path=None, + advantage_estimator=estimator, + rewards_normalization=True, + grpo_std_normalization=grpo_std_normalization, + n_samples_per_prompt=n_samples_per_prompt, + reward_key=None, + ) + samples = [ + Sample(group_index=position // n_samples_per_prompt, reward=reward) for position, reward in enumerate(rewards) + ] + return post_process_rewards(args, samples) + + +def test_p3o_registry_uses_feynrl_sample_std_independent_of_grpo_flag(): + p3o_raw, p3o_normalized = _normalized_rewards("p3o") + _, p3o_without_grpo_flag = _normalized_rewards("p3o", grpo_std_normalization=False) + + assert p3o_raw == [1.0, 3.0, 2.0, 6.0] + expected = [-1 / math.sqrt(2), 1 / math.sqrt(2)] * 2 + assert p3o_normalized == pytest.approx(expected, abs=1e-6) + assert p3o_without_grpo_flag == pytest.approx(expected, abs=1e-6) + + +def test_p3o_registry_preserves_raw_reward_for_single_sample_groups(): + raw, normalized = _normalized_rewards( + "p3o", + rewards=(1.5, -2.0), + n_samples_per_prompt=1, + ) + + assert raw == normalized == [1.5, -2.0] + + +def test_grpo_registry_normalization_is_unchanged(): + _, normalized = _normalized_rewards("grpo", grpo_std_normalization=False) + + assert normalized == [-1.0, 1.0, -2.0, 2.0] diff --git a/tests/utils/test_rollout_logprob_mask.py b/tests/utils/test_rollout_logprob_mask.py new file mode 100644 index 000000000..e6213f4ee --- /dev/null +++ b/tests/utils/test_rollout_logprob_mask.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Behavior-logprob pairing-mask tests.""" + +from types import SimpleNamespace + +import pytest + +from relax.utils.training.data_fields import build_data_fields +from relax.utils.types import Sample +from relax.utils.utils import convert_samples_to_train_data + + +def _args(**overrides): + values = { + "advantage_estimator": "p3o", + "agentic_custom_advantage_path": None, + "custom_reward_post_process_path": None, + "debug_train_only": True, + "grpo_std_normalization": True, + "loss_type": "policy_loss", + "multimodal_keys": None, + "n_samples_per_prompt": 1, + "reward_key": None, + "rewards_normalization": False, + "use_opd": False, + "use_rollout_logprobs": True, + "use_rollout_routing_replay": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _sample(*, pairing_mask=None, rollout_log_probs=None, loss_mask=None): + return Sample( + tokens=[100, 1, 2, 3], + response_length=3, + reward=1.0, + loss_mask=loss_mask, + rollout_log_probs=rollout_log_probs, + rollout_log_probs_mask=pairing_mask, + ) + + +def test_pairing_mask_is_carried_and_intersected_with_loss_mask() -> None: + sample = _sample( + pairing_mask=[True, False, True], + rollout_log_probs=[-0.1, -0.2, -0.3], + loss_mask=[1, 1, 0], + ) + + train_data = convert_samples_to_train_data(_args(), [sample]) + + assert train_data["rollout_log_probs_mask"] == [[True, False, True]] + assert train_data["loss_masks"] == [[1, 0, 0]] + assert "rollout_log_probs_mask" in build_data_fields(_args()) + + +def test_missing_pairing_mask_defaults_to_all_true_without_changing_loss_mask() -> None: + sample = _sample(rollout_log_probs=[-0.1, -0.2, -0.3], loss_mask=[1, 0, 1]) + + train_data = convert_samples_to_train_data(_args(), [sample]) + + assert train_data["rollout_log_probs_mask"] == [[True, True, True]] + assert train_data["loss_masks"] == [[1, 0, 1]] + + +@pytest.mark.parametrize( + ("rollout_log_probs", "pairing_mask", "match"), + [ + ([-0.1, -0.2], None, "rollout log-prob length"), + ([-0.1, -0.2, -0.3], [True, False], "rollout log-prob mask length"), + ], +) +def test_pairing_alignment_mismatch_is_rejected(rollout_log_probs, pairing_mask, match) -> None: + sample = _sample(rollout_log_probs=rollout_log_probs, pairing_mask=pairing_mask) + + with pytest.raises(ValueError, match=match): + convert_samples_to_train_data(_args(), [sample]) + + +def test_requested_behavior_logprobs_cannot_be_missing() -> None: + with pytest.raises(ValueError, match="requires behavior log-probs"): + convert_samples_to_train_data(_args(), [_sample()]) diff --git a/tests/utils/training/test_p3o_replay.py b/tests/utils/training/test_p3o_replay.py new file mode 100644 index 000000000..5bec7f867 --- /dev/null +++ b/tests/utils/training/test_p3o_replay.py @@ -0,0 +1,194 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for P3O's two-pass replay guards and stat-accumulation scope. + +The pieces under test here are the ones that decide *which tokens* enter ESS and +*whether the window can be replayed* -- the two places where a wrong answer still +produces a plausible-looking loss curve. The distributed matrix (DP/CP/TP/PP) and +the end-to-end training run require multi-GPU and are covered separately. +""" + +import sys +from types import ModuleType + +import pytest +import torch + +from relax.utils.training.p3o_replay import ( + preserved_iterator_positions, + preserved_rng_state, +) +from relax.utils.training.p3o_utils import ( + P3OSufficientStats, + finalize_p3o_step_context, +) + + +TOL = dict(rel=1e-6, abs=1e-6) + + +class _FakeIterator: + """Minimal stand-in exposing the replay contract used by the pre-pass.""" + + def __init__(self, items): + self.items = list(items) + self.offset = 0 + + def __next__(self): + if self.offset >= len(self.items): + raise StopIteration + item = self.items[self.offset] + self.offset += 1 + return item + + def snapshot_position(self) -> int: + return self.offset + + def restore_position(self, position: int) -> None: + self.offset = position + + +def test_p3o_iterator_positions_restored_after_prepass(): + iterator = _FakeIterator(range(6)) + next(iterator) + next(iterator) + assert iterator.offset == 2 + + with preserved_iterator_positions([iterator]): + next(iterator) + next(iterator) + assert iterator.offset == 4 + + # Restores to mid-rollout position, not to zero. + assert iterator.offset == 2 + + +def test_p3o_iterator_positions_restored_even_when_prepass_raises(): + iterator = _FakeIterator(range(6)) + next(iterator) + + with pytest.raises(RuntimeError, match="boom"): + with preserved_iterator_positions([iterator]): + next(iterator) + raise RuntimeError("boom") + + assert iterator.offset == 1 + + +def test_p3o_duplicate_iterator_instances_restored_once(): + """Virtual PP passes the same iterator once per model chunk.""" + iterator = _FakeIterator(range(6)) + next(iterator) + + with preserved_iterator_positions([iterator, iterator, None]): + next(iterator) + + assert iterator.offset == 1 + + +def test_p3o_non_replayable_iterator_is_rejected_loudly(): + class _Opaque: + pass + + with pytest.raises(RuntimeError, match="not replayable"): + with preserved_iterator_positions([_Opaque()]): + pass + + +def test_p3o_rng_state_restored_after_prepass(): + torch.manual_seed(1234) + expected = torch.randn(4) + + torch.manual_seed(1234) + with preserved_rng_state(): + # Burn RNG inside the pre-pass, as a stochastic forward would. + torch.randn(16) + actual = torch.randn(4) + + torch.testing.assert_close(actual, expected) + + +def test_p3o_rng_and_megatron_tracker_restored_after_error(monkeypatch): + # preserved_rng_state() imports the tracker lazily from megatron. CI installs no + # megatron, so supply just the one module that import needs; a real install is + # used as-is, keeping the GPU path identical. + megatron_random = sys.modules.get("megatron.core.tensor_parallel.random") + if megatron_random is None: + for name in ( + "megatron", + "megatron.core", + "megatron.core.tensor_parallel", + "megatron.core.tensor_parallel.random", + ): + module = ModuleType(name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, name, module) + megatron_random = sys.modules["megatron.core.tensor_parallel.random"] + # Seed the symbol so the monkeypatch below patches rather than invents it, + # matching how a real megatron module would look at import time. + megatron_random.get_cuda_rng_tracker = lambda: None + + class _FakeTracker: + def __init__(self): + self.states = {"model-parallel-rng": torch.tensor([7], dtype=torch.uint8)} + + def get_states(self): + return {name: state.clone() for name, state in self.states.items()} + + def set_states(self, states): + self.states = {name: state.clone() for name, state in states.items()} + + tracker = _FakeTracker() + monkeypatch.setattr(megatron_random, "get_cuda_rng_tracker", lambda: tracker) + + torch.manual_seed(2026) + expected = torch.randn(4) + torch.manual_seed(2026) + + with pytest.raises(RuntimeError, match="stats pass failed"): + with preserved_rng_state(): + torch.randn(8) + tracker.states["model-parallel-rng"] = torch.tensor([99], dtype=torch.uint8) + raise RuntimeError("stats pass failed") + + torch.testing.assert_close(torch.randn(4), expected) + torch.testing.assert_close( + tracker.states["model-parallel-rng"], + torch.tensor([7], dtype=torch.uint8), + ) + + +def test_p3o_stats_accumulate_then_reduce_equals_single_shot(): + """Sum-then-reduce must equal computing over the concatenated token set.""" + shards = [ + P3OSufficientStats( + sum_ratio=torch.tensor(1.5, dtype=torch.float64), + sum_ratio_sq=torch.tensor(2.25, dtype=torch.float64), + valid_token_count=torch.tensor(1.0, dtype=torch.float64), + ), + P3OSufficientStats( + sum_ratio=torch.tensor(6.0, dtype=torch.float64), + sum_ratio_sq=torch.tensor(19.0, dtype=torch.float64), + valid_token_count=torch.tensor(3.0, dtype=torch.float64), + ), + ] + total = shards[0] + shards[1] + + assert float(total.sum_ratio) == pytest.approx(7.5, **TOL) + assert float(total.sum_ratio_sq) == pytest.approx(21.25, **TOL) + assert float(total.valid_token_count) == 4.0 + assert float(finalize_p3o_step_context(total).normalized_ess) == pytest.approx(0.6617647055709343, **TOL) + + +def test_p3o_dummy_microbatch_contributes_nothing(): + """Dummy micro-batches align DP counts and must not move ESS.""" + real = P3OSufficientStats( + sum_ratio=torch.tensor(7.5, dtype=torch.float64), + sum_ratio_sq=torch.tensor(21.25, dtype=torch.float64), + valid_token_count=torch.tensor(4.0, dtype=torch.float64), + ) + with_dummy = real + P3OSufficientStats.zeros() + + assert float(finalize_p3o_step_context(with_dummy).normalized_ess) == pytest.approx( + float(finalize_p3o_step_context(real).normalized_ess), **TOL + ) diff --git a/tests/utils/training/test_p3o_utils.py b/tests/utils/training/test_p3o_utils.py new file mode 100644 index 000000000..ebbd87820 --- /dev/null +++ b/tests/utils/training/test_p3o_utils.py @@ -0,0 +1,501 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Element-wise parity tests for the P3O primitives. + +The golden values come from running the reference implementation (FeynRL +``algs/P3O/p3o.py``) over one logical batch. The same element-wise oracle is +used by the default micro-batch scope and the optional optimizer-step scope. +""" + +import math + +import pytest +import torch + +from relax.utils.training.p3o_utils import ( + P3OStepContext, + P3OSufficientStats, + compute_p3o_behavior_kl_proxy, + compute_p3o_exact_kl, + compute_p3o_sufficient_stats, + compute_p3o_token_terms, + finalize_p3o_step_context, +) + + +# Golden case: ratios [1.0, 2.0, 0.5, 4.0] laid out as two sequences of three +# tokens each, with the third token of every sequence invalid (padding). +GOLDEN_RATIOS = [1.0, 2.0, 0.5, 4.0] +GOLDEN_S1 = 7.5 +GOLDEN_S2 = 21.25 +GOLDEN_N = 4 +GOLDEN_ESS = 0.6617647055709343 +GOLDEN_LOSS_MEAN = 0.8332794905 +GOLDEN_GRAD = [ + [-0.6617646813, 0.8308823705, 0.0], + [-1.3382353783, 0.5845587850, 0.0], +] + +# pytest.approx uses rel/abs; torch.testing.assert_close uses rtol/atol. +TOL = dict(rel=1e-6, abs=1e-6) +TENSOR_TOL = dict(rtol=1e-6, atol=1e-6) + + +GOLDEN_BEHAVIOR_LOG_PROB = -2.0 +GOLDEN_ADVANTAGES = [[1.0, -1.0, 0.0], [2.0, -0.5, 0.0]] +GOLDEN_COEFFICIENTS = [[GOLDEN_ESS, GOLDEN_ESS, 0.0], [0.5, GOLDEN_ESS, 0.0]] +GOLDEN_TOKEN_TOTALS = [ + [1.3235294111, -0.7994998778, 0.0], + [2.7969356343, 0.0121528449, 0.0], +] + + +def _golden_batch(requires_grad: bool = False): + """Build the golden 2x3 batch: ratios above, pad in column 2. + + The behavior log-prob level and the advantages are part of the frozen golden + case: the loss value pins the log-prob level (the score term is + ``-coef * log_prob * A``), while the four gradients pin the advantages. + """ + behavior_log_probs = torch.full((2, 3), GOLDEN_BEHAVIOR_LOG_PROB, dtype=torch.float32) + log_ratio = torch.tensor( + [[math.log(1.0), math.log(2.0), 0.0], [math.log(0.5), math.log(4.0), 0.0]], + dtype=torch.float32, + ) + log_probs = (behavior_log_probs + log_ratio).clone() + log_probs.requires_grad_(requires_grad) + advantages = torch.tensor(GOLDEN_ADVANTAGES, dtype=torch.float32) + valid_mask = torch.tensor([[True, True, False], [True, True, False]]) + return log_probs, behavior_log_probs, advantages, valid_mask + + +def _mean_loss(terms, valid_mask): + """Token-sum of the full objective normalized by the global valid count.""" + total = (terms.score_loss + terms.adaptive_kl_loss).sum() + return total / valid_mask.sum() + + +def test_p3o_utils_sufficient_stats_match_reference_moments(): + log_probs, behavior_log_probs, _, valid_mask = _golden_batch() + stats = compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + + assert float(stats.sum_ratio) == pytest.approx(GOLDEN_S1, **TOL) + assert float(stats.sum_ratio_sq) == pytest.approx(GOLDEN_S2, **TOL) + assert float(stats.valid_token_count) == GOLDEN_N + + +def test_p3o_utils_normalized_ess_matches_reference(): + stats = P3OSufficientStats( + sum_ratio=torch.tensor(GOLDEN_S1, dtype=torch.float64), + sum_ratio_sq=torch.tensor(GOLDEN_S2, dtype=torch.float64), + valid_token_count=torch.tensor(float(GOLDEN_N), dtype=torch.float64), + ) + ctx = finalize_p3o_step_context(stats) + + assert float(ctx.normalized_ess) == pytest.approx(GOLDEN_ESS, **TOL) + assert float(ctx.adaptive_cap) == pytest.approx(GOLDEN_ESS, **TOL) + assert float(ctx.valid_token_count) == GOLDEN_N + assert float(ctx.ratio_mean) == pytest.approx(GOLDEN_S1 / GOLDEN_N, **TOL) + assert ctx.clamp_events == 0 + + +def test_p3o_utils_total_loss_matches_reference_golden_value(): + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch() + stats = compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + ctx = finalize_p3o_step_context(stats) + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + + assert float(_mean_loss(terms, valid_mask)) == pytest.approx(GOLDEN_LOSS_MEAN, **TOL) + + +def test_p3o_reference_oracle_matches_ess_cap_and_token_loss(): + """Expose the complete FeynRL formula oracle in one elementwise check.""" + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch() + context = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, context) + + expected_coefficients = torch.tensor(GOLDEN_COEFFICIENTS, dtype=torch.float32) + expected_token_totals = torch.tensor(GOLDEN_TOKEN_TOTALS, dtype=torch.float32) + + assert float(context.normalized_ess) == pytest.approx(GOLDEN_ESS, **TOL) + assert float(context.adaptive_cap) == pytest.approx(GOLDEN_ESS, **TOL) + torch.testing.assert_close( + torch.minimum(terms.ratio, context.adaptive_cap.float()), + expected_coefficients, + **TENSOR_TOL, + ) + torch.testing.assert_close( + terms.score_loss + terms.adaptive_kl_loss, + expected_token_totals, + **TENSOR_TOL, + ) + assert float(_mean_loss(terms, valid_mask)) == pytest.approx(GOLDEN_LOSS_MEAN, **TOL) + + +def test_p3o_utils_gradient_matches_reference_golden_value(): + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch(requires_grad=True) + stats = compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + ctx = finalize_p3o_step_context(stats) + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + + (terms.score_loss + terms.adaptive_kl_loss).sum().backward() + + expected = torch.tensor(GOLDEN_GRAD, dtype=torch.float32) + torch.testing.assert_close(log_probs.grad, expected, **TENSOR_TOL) + + +def test_p3o_utils_ess_invariant_to_token_partitioning(): + """Splitting the same tokens across micro-batches must not move the cap.""" + log_probs, behavior_log_probs, _, valid_mask = _golden_batch() + + whole = compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + accumulated = P3OSufficientStats.zeros() + for row in range(log_probs.shape[0]): + accumulated = accumulated + compute_p3o_sufficient_stats( + log_probs[row : row + 1], behavior_log_probs[row : row + 1], valid_mask[row : row + 1] + ) + + whole_ess = float(finalize_p3o_step_context(whole).normalized_ess) + split_ess = float(finalize_p3o_step_context(accumulated).normalized_ess) + assert whole_ess == pytest.approx(split_ess, **TOL) + assert whole_ess == pytest.approx(GOLDEN_ESS, **TOL) + + +def test_p3o_utils_dp_cp_stat_reduction_matches_single_rank(): + """Per-rank shards summed elementwise reproduce the single-rank moments.""" + rank0 = P3OSufficientStats( + sum_ratio=torch.tensor(3.0, dtype=torch.float64), + sum_ratio_sq=torch.tensor(5.0, dtype=torch.float64), + valid_token_count=torch.tensor(2.0, dtype=torch.float64), + ) + rank1 = P3OSufficientStats( + sum_ratio=torch.tensor(4.5, dtype=torch.float64), + sum_ratio_sq=torch.tensor(16.25, dtype=torch.float64), + valid_token_count=torch.tensor(2.0, dtype=torch.float64), + ) + reduced = P3OSufficientStats.from_vector(rank0.as_vector() + rank1.as_vector()) + + assert float(reduced.sum_ratio) == pytest.approx(GOLDEN_S1, **TOL) + assert float(reduced.sum_ratio_sq) == pytest.approx(GOLDEN_S2, **TOL) + assert float(reduced.valid_token_count) == GOLDEN_N + assert float(finalize_p3o_step_context(reduced).normalized_ess) == pytest.approx(GOLDEN_ESS, **TOL) + + +def test_p3o_utils_on_policy_degenerates_to_vanilla_policy_gradient(): + """rho == 1 everywhere => cap == 1, adaptive KL == 0, gradient == PG.""" + behavior_log_probs = torch.full((2, 4), -0.5, dtype=torch.float32) + log_probs = behavior_log_probs.clone().requires_grad_(True) + advantages = torch.tensor([[1.0, -2.0, 0.5, 1.5], [-1.0, 2.0, -0.5, 0.25]], dtype=torch.float32) + valid_mask = torch.ones(2, 4, dtype=torch.bool) + + ctx = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + assert float(ctx.normalized_ess) == pytest.approx(1.0, **TOL) + + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + torch.testing.assert_close(terms.adaptive_kl_loss, torch.zeros_like(terms.adaptive_kl_loss), **TENSOR_TOL) + + (terms.score_loss + terms.adaptive_kl_loss).sum().backward() + torch.testing.assert_close(log_probs.grad, -advantages, **TENSOR_TOL) + + +def test_p3o_utils_uniform_ratio_offset_leaves_ess_near_one(): + """ESS measures concentration, so a constant shift is not mismatch.""" + behavior_log_probs = torch.zeros(2, 4, dtype=torch.float32) + log_probs = behavior_log_probs + 0.75 + valid_mask = torch.ones(2, 4, dtype=torch.bool) + + ctx = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + assert float(ctx.normalized_ess) == pytest.approx(1.0, **TOL) + + +def test_p3o_utils_dominant_ratio_drives_ess_toward_one_over_n(): + """One huge ratio among N tokens collapses ESS to roughly 1/N.""" + behavior_log_probs = torch.zeros(1, 4, dtype=torch.float32) + log_probs = torch.tensor([[math.log(1e6), 0.0, 0.0, 0.0]], dtype=torch.float32) + valid_mask = torch.ones(1, 4, dtype=torch.bool) + + ctx = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + assert float(ctx.normalized_ess) == pytest.approx(0.25, rel=1e-3) + + +def test_p3o_utils_single_valid_token_gives_full_ess(): + behavior_log_probs = torch.zeros(1, 3, dtype=torch.float32) + log_probs = torch.tensor([[math.log(3.0), 0.0, 0.0]], dtype=torch.float32) + valid_mask = torch.tensor([[True, False, False]]) + + ctx = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + assert float(ctx.normalized_ess) == pytest.approx(1.0, **TOL) + assert float(ctx.valid_token_count) == 1 + + +def test_p3o_utils_masked_positions_tolerate_non_finite_values(): + """NaN/Inf in prompt or padding slots must not leak into the stats.""" + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch() + log_probs, behavior_log_probs = log_probs.clone(), behavior_log_probs.clone() + advantages = advantages.clone() + for tensor, poison in ((log_probs, float("nan")), (behavior_log_probs, float("inf")), (advantages, 1e30)): + tensor[0, 2] = poison + tensor[1, 2] = -poison if poison == 1e30 else float("nan") + + stats = compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + ctx = finalize_p3o_step_context(stats) + assert float(ctx.normalized_ess) == pytest.approx(GOLDEN_ESS, **TOL) + + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + assert torch.isfinite(terms.score_loss).all() + assert float(_mean_loss(terms, valid_mask)) == pytest.approx(GOLDEN_LOSS_MEAN, **TOL) + + +@pytest.mark.parametrize( + ("log_prob", "behavior_log_prob"), + [ + (float("nan"), 0.0), + (float("inf"), 0.0), + (float("-inf"), 0.0), + (0.0, float("inf")), + (0.0, float("-inf")), + ], +) +def test_p3o_utils_non_finite_valid_token_raises(log_prob, behavior_log_prob): + behavior_log_probs = torch.tensor([[behavior_log_prob, 0.0]], dtype=torch.float32) + log_probs = torch.tensor([[log_prob, 0.0]], dtype=torch.float32) + valid_mask = torch.ones(1, 2, dtype=torch.bool) + + with pytest.raises(ValueError, match="non-finite importance ratio"): + compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + + +def test_p3o_utils_all_masked_poison_produces_fp64_zero_stats(): + log_probs = torch.tensor([[float("nan"), float("inf")]], dtype=torch.float32) + behavior_log_probs = torch.tensor([[float("-inf"), float("nan")]], dtype=torch.float32) + valid_mask = torch.zeros(1, 2, dtype=torch.bool) + + stats = compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + + for value in (stats.sum_ratio, stats.sum_ratio_sq, stats.valid_token_count): + assert value.dtype == torch.float64 + assert torch.equal(value, torch.zeros((), dtype=torch.float64)) + + +def test_p3o_utils_empty_global_batch_falls_back_to_full_ess(): + stats = P3OSufficientStats.zeros() + context = finalize_p3o_step_context(stats) + + assert float(context.normalized_ess) == 1.0 + assert float(context.adaptive_cap) == 1.0 + assert float(context.valid_token_count) == 0.0 + assert float(context.ratio_mean) == 1.0 + assert float(context.ratio_std) == 0.0 + + +@pytest.mark.parametrize("mismatched", ["behavior", "mask"]) +def test_p3o_utils_sufficient_stats_reject_shape_mismatch(mismatched): + log_probs = torch.zeros(2, 3) + behavior_log_probs = torch.zeros(2, 2) if mismatched == "behavior" else torch.zeros(2, 3) + valid_mask = torch.ones(2, 2, dtype=torch.bool) if mismatched == "mask" else torch.ones(2, 3, dtype=torch.bool) + + with pytest.raises(ValueError, match="identical shapes"): + compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + + +def test_p3o_utils_token_terms_reject_advantage_shape_mismatch(): + log_probs, behavior_log_probs, _, valid_mask = _golden_batch() + context = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + + with pytest.raises(ValueError, match="advantages"): + compute_p3o_token_terms( + log_probs, + behavior_log_probs, + torch.zeros(2, 1), + valid_mask, + context, + ) + + +def test_p3o_utils_cap_hits_track_ratios_above_cap(): + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch() + ctx = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + + # ratios 1.0, 2.0, 4.0 exceed cap 0.6617...; ratio 0.5 does not; pads never count. + expected = torch.tensor([[1.0, 1.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float32) + torch.testing.assert_close(terms.cap_hits, expected) + assert float(terms.cap_hits.sum() / ctx.valid_token_count) == pytest.approx(0.75, **TOL) + + +def test_p3o_utils_behavior_kl_proxy_is_non_negative_and_directional(): + behavior_log_probs = torch.zeros(1, 3, dtype=torch.float32) + log_probs = torch.tensor([[math.log(2.0), math.log(0.5), 0.0]], dtype=torch.float32) + valid_mask = torch.ones(1, 3, dtype=torch.bool) + + kl = compute_p3o_behavior_kl_proxy(log_probs, behavior_log_probs, valid_mask) + + assert (kl >= -1e-7).all() + assert float(kl[0, 2]) == pytest.approx(0.0, abs=1e-7) + # k3 form: l + exp(-l) - 1 + assert float(kl[0, 0]) == pytest.approx(math.log(2.0) + 0.5 - 1.0, **TOL) + assert float(kl[0, 1]) == pytest.approx(math.log(0.5) + 2.0 - 1.0, **TOL) + + +def test_p3o_utils_behavior_kl_proxy_clamps_extreme_divergence(): + behavior_log_probs = torch.zeros(1, 1, dtype=torch.float32) + log_probs = torch.tensor([[-50.0]], dtype=torch.float32) + valid_mask = torch.ones(1, 1, dtype=torch.bool) + + kl = compute_p3o_behavior_kl_proxy(log_probs, behavior_log_probs, valid_mask) + assert float(kl[0, 0]) == pytest.approx(-50.0 + math.exp(10.0) - 1.0, rel=1e-6) + + +def test_p3o_utils_proxy_safe_matches_proxy_forward_and_has_correct_gradient_sign(): + behavior_log_probs = torch.zeros(121, dtype=torch.float32) + proxy_log_probs = torch.linspace(-30.0, 30.0, 121, requires_grad=True) + safe_log_probs = proxy_log_probs.detach().clone().requires_grad_(True) + valid_mask = torch.ones_like(proxy_log_probs, dtype=torch.bool) + + proxy = compute_p3o_behavior_kl_proxy(proxy_log_probs, behavior_log_probs, valid_mask, mode="proxy") + proxy_safe = compute_p3o_behavior_kl_proxy(safe_log_probs, behavior_log_probs, valid_mask, mode="proxy_safe") + + torch.testing.assert_close(proxy_safe, proxy, rtol=0.0, atol=0.0) + proxy.sum().backward() + proxy_safe.sum().backward() + + negative = safe_log_probs.detach() < 0 + positive = safe_log_probs.detach() > 0 + assert torch.all(safe_log_probs.grad[negative] <= 0) + assert torch.all(safe_log_probs.grad[positive] >= 0) + assert float(safe_log_probs.grad.abs().max()) <= math.exp(10.0) + assert float(proxy_log_probs.grad[0]) > 0 + assert float(safe_log_probs.grad[0]) < 0 + + +def test_p3o_utils_exact_kl_matches_manual_small_vocabulary_oracle(): + policy_logits = torch.tensor([[[1.0, 0.0, -1.0], [float("nan"), 2.0, 1.0]]], requires_grad=True) + behavior_logits = torch.tensor([[[0.0, 0.5, -0.5], [float("inf"), 0.0, 0.0]]], requires_grad=True) + valid_mask = torch.tensor([[True, False]]) + + exact = compute_p3o_exact_kl(policy_logits, behavior_logits, valid_mask) + policy_log_probs = torch.log_softmax(policy_logits[0, 0], dim=-1) + behavior_log_probs = torch.log_softmax(behavior_logits[0, 0].detach(), dim=-1) + expected = (policy_log_probs.exp() * (policy_log_probs - behavior_log_probs)).sum() + + torch.testing.assert_close(exact[0, 0], expected) + assert float(exact[0, 1].detach()) == 0.0 + exact.sum().backward() + assert policy_logits.grad is not None + assert behavior_logits.grad is None + + +def test_p3o_utils_exact_training_mode_requires_behavior_logits(): + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch() + context = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + + with pytest.raises(ValueError, match="full-vocabulary behavior logits"): + compute_p3o_token_terms( + log_probs, + behavior_log_probs, + advantages, + valid_mask, + context, + kl_mode="exact", + ) + + +def test_p3o_utils_advantage_and_cap_are_stop_gradient(): + log_probs, behavior_log_probs, advantages, valid_mask = _golden_batch(requires_grad=True) + advantages = advantages.clone().requires_grad_(True) + behavior_log_probs = behavior_log_probs.clone().requires_grad_(True) + + ctx = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + (terms.score_loss + terms.adaptive_kl_loss).sum().backward() + + assert advantages.grad is None + assert behavior_log_probs.grad is None + assert log_probs.grad is not None + assert not ctx.normalized_ess.requires_grad + + +def test_p3o_utils_entire_adaptive_coefficient_is_stop_gradient(): + log_probs = torch.tensor([math.log(2.0)], dtype=torch.float32, requires_grad=True) + behavior_log_probs = torch.zeros(1, dtype=torch.float32) + advantages = torch.tensor([2.0], dtype=torch.float32) + valid_mask = torch.ones(1, dtype=torch.bool) + adaptive_cap = torch.tensor(0.75, dtype=torch.float64, requires_grad=True) + ctx = finalize_p3o_step_context( + P3OSufficientStats( + sum_ratio=torch.tensor(1.0, dtype=torch.float64), + sum_ratio_sq=torch.tensor(1.0, dtype=torch.float64), + valid_token_count=torch.tensor(1.0, dtype=torch.float64), + ) + ) + ctx = type(ctx)( + normalized_ess=ctx.normalized_ess, + adaptive_cap=adaptive_cap, + valid_token_count=ctx.valid_token_count, + ratio_mean=ctx.ratio_mean, + ratio_std=ctx.ratio_std, + ) + + terms = compute_p3o_token_terms(log_probs, behavior_log_probs, advantages, valid_mask, ctx) + terms.score_loss.sum().backward() + + torch.testing.assert_close(log_probs.grad, torch.tensor([-1.5])) + assert adaptive_cap.grad is None + + +def test_p3o_utils_clip_hits_use_monitoring_interval_not_adaptive_cap(): + behavior_log_probs = torch.zeros(1, 5) + ratios = torch.tensor([[0.79, 0.8, 1.0, 1.2, 1.21]]) + log_probs = ratios.log() + valid_mask = torch.ones_like(log_probs, dtype=torch.bool) + context = finalize_p3o_step_context(compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask)) + + terms = compute_p3o_token_terms( + log_probs, + behavior_log_probs, + torch.ones_like(log_probs), + valid_mask, + context, + clip_low=0.2, + clip_high=0.2, + ) + + torch.testing.assert_close(terms.clip_hits, torch.tensor([[1.0, 0.0, 0.0, 0.0, 1.0]])) + + +def test_p3o_utils_token_terms_keep_adaptive_cap_on_device(monkeypatch): + """The per-micro-batch loss must not convert the GPU cap to a scalar.""" + adaptive_cap = torch.tensor(0.75, dtype=torch.float64) + context = P3OStepContext( + normalized_ess=adaptive_cap, + adaptive_cap=adaptive_cap, + valid_token_count=torch.tensor(1.0, dtype=torch.float64), + ratio_mean=torch.tensor(2.0, dtype=torch.float64), + ratio_std=torch.tensor(0.0, dtype=torch.float64), + ) + log_probs = torch.tensor([math.log(2.0)], dtype=torch.float32, requires_grad=True) + + def fail_on_scalar_conversion(tensor): + raise AssertionError(f"unexpected Tensor.__float__ for {tensor}") + + monkeypatch.setattr(torch.Tensor, "__float__", fail_on_scalar_conversion) + terms = compute_p3o_token_terms( + log_probs=log_probs, + behavior_log_probs=torch.zeros_like(log_probs), + advantages=torch.ones_like(log_probs), + valid_mask=torch.ones_like(log_probs, dtype=torch.bool), + step_context=context, + ) + + torch.testing.assert_close(terms.score_loss, -adaptive_cap.float() * log_probs.detach()) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.bfloat16]) +def test_p3o_utils_stats_stable_across_input_dtypes(dtype): + log_probs, behavior_log_probs, _, valid_mask = _golden_batch() + stats = compute_p3o_sufficient_stats(log_probs.to(dtype), behavior_log_probs.to(dtype), valid_mask) + ess = float(finalize_p3o_step_context(stats).normalized_ess) + + assert stats.as_vector().dtype == torch.float64 + tol = 5e-3 if dtype is torch.bfloat16 else 1e-6 + assert ess == pytest.approx(GOLDEN_ESS, rel=tol, abs=tol) From c007b8ee77847a9a3f89c6f4672b2f1a3e98d682 Mon Sep 17 00:00:00 2001 From: Chream Date: Thu, 13 Aug 2026 18:56:43 +0800 Subject: [PATCH 02/30] refactor(megatron): clarify CP merge metadata --- relax/backends/megatron/cp_utils.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/relax/backends/megatron/cp_utils.py b/relax/backends/megatron/cp_utils.py index 542994d16..1e75d316c 100644 --- a/relax/backends/megatron/cp_utils.py +++ b/relax/backends/megatron/cp_utils.py @@ -649,24 +649,32 @@ def dynamic_cp_merge_output( # 1. reconstruct each sample's full response from its CP-local zig-zag shards. if dynamic_cp_size > 1: dynamic_cp_group = mpu.get_dynamic_data_context_parallel_groups(group_size=dynamic_cp_size) - ptls = padded_total_lengths if padded_total_lengths is not None else [None] * len(values) + padded_total_lengths_for_values = ( + padded_total_lengths if padded_total_lengths is not None else [None] * len(values) + ) _validate_metadata_lengths( values=values, total_lengths=total_lengths, response_lengths=response_lengths, - padded_total_lengths=ptls, + padded_total_lengths=padded_total_lengths_for_values, ) values = [ all_gather_with_cp( - v, - tl, - rl, - padded_total_length=ptl, + value, + total_length, + response_length, + padded_total_length=padded_total_length, dynamic_cp_size=dynamic_cp_size, dynamic_cp_rank=dynamic_cp_rank, dynamic_cp_group=dynamic_cp_group, ) - for v, tl, rl, ptl in zip(values, total_lengths, response_lengths, ptls, strict=True) + for value, total_length, response_length, padded_total_length in zip( + values, + total_lengths, + response_lengths, + padded_total_lengths_for_values, + strict=True, + ) ] # 2. collect all sub-groups' samples across the static CP group and reorder. From 3313cf569407ad3f75ec29174e2c0c4814c4386e Mon Sep 17 00:00:00 2001 From: Chream Date: Thu, 13 Aug 2026 18:56:53 +0800 Subject: [PATCH 03/30] docs(p3o): format Chinese recipe table --- examples/algorithms/p3o/README_zh.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md index fc1fb53cf..ad26218a2 100644 --- a/examples/algorithms/p3o/README_zh.md +++ b/examples/algorithms/p3o/README_zh.md @@ -101,12 +101,12 @@ reduction 分块。这两个设置对 P3O 和 GRPO 完全一致，都会记录 ## 场景 -| 场景 | 更新间隔 | 温度 override | 含义 | -| --- | ---: | ---: | --- | -| `on_policy` | 1 | 关闭 | 以正常采样配置在每个 rollout 后同步。 | -| `periodic_sync_interval_3` | 3 | 关闭 | 只引入周期性的 rollout-policy 陈旧性。 | -| `temperature_0p6` | 1 | 0.6 | 只改变行为策略温度。 | -| `temperature_1p2` | 1 | 1.2 | 只改变行为策略温度。 | +| 场景 | 更新间隔 | 温度 override | 含义 | +| -------------------------- | -------: | ------------: | -------------------------------------- | +| `on_policy` | 1 | 关闭 | 以正常采样配置在每个 rollout 后同步。 | +| `periodic_sync_interval_3` | 3 | 关闭 | 只引入周期性的 rollout-policy 陈旧性。 | +| `temperature_0p6` | 1 | 0.6 | 只改变行为策略温度。 | +| `temperature_1p2` | 1 | 1.2 | 只改变行为策略温度。 | 同一场景中的 P3O 和 GRPO 启动器共享所有非算法配置。温度场景会保持 `top_p`、 `top_k`、response 限制和评估采样设置不变。 From 22635d4fc36661e8665d6ab595321705c5e16fdb Mon Sep 17 00:00:00 2001 From: Chream Date: Fri, 14 Aug 2026 10:32:21 +0800 Subject: [PATCH 04/30] fix(p3o): harden behavior policy safeguards --- examples/algorithms/p3o/rollout.py | 26 ++++ relax/backends/megatron/actor.py | 15 +- relax/backends/megatron/checkpoint.py | 10 +- relax/backends/megatron/loss.py | 35 ++++- relax/backends/megatron/model.py | 88 ++++++++++- relax/backends/megatron/rollout_policy_lag.py | 46 +++++- relax/engine/rollout/sglang_rollout.py | 47 ++++++ relax/utils/arguments.py | 93 ++++++++---- .../backends/megatron/test_p3o_distributed.py | 143 +++++++++++++++++- tests/backends/megatron/test_p3o_loss.py | 41 ++++- .../backends/megatron/test_p3o_model_step.py | 95 +++++++++++- .../megatron/test_p3o_observability.py | 129 ++++++++++++++++ .../rollout/test_p3o_sampling_contract.py | 73 +++++++++ tests/examples/algorithms/p3o/test_rollout.py | 26 +++- tests/utils/test_p3o_arguments.py | 97 +++++++++++- 15 files changed, 903 insertions(+), 61 deletions(-) create mode 100644 tests/engine/rollout/test_p3o_sampling_contract.py diff --git a/examples/algorithms/p3o/rollout.py b/examples/algorithms/p3o/rollout.py index a51e7fb61..eca88ebd4 100644 --- a/examples/algorithms/p3o/rollout.py +++ b/examples/algorithms/p3o/rollout.py @@ -10,6 +10,15 @@ from relax.utils.types import Sample +_P3O_TRUNCATION_SAMPLING_KEYS = ( + "min_p", + "top_a", + "typical_p", + "epsilon_cutoff", + "eta_cutoff", +) + + async def _sglang_generate(*args: Any, **kwargs: Any) -> Sample: """Import the heavyweight rollout backend only when generation starts.""" from relax.engine.rollout.sglang_rollout import generate @@ -34,6 +43,17 @@ def behavior_sampling_params(sampling_params: dict[str, Any], *, evaluation: boo """Return isolated sampling parameters for P3O rollout generation.""" updated = sampling_params.copy() if not evaluation: + if updated.get("top_p", 1.0) != 1.0 or updated.get("top_k", -1) != -1: + raise ValueError( + "P3O behavior sampling requires top_p=1.0 and top_k=-1 so rollout log-probs describe " + "the untruncated distribution" + ) + unsupported = sorted(key for key in _P3O_TRUNCATION_SAMPLING_KEYS if key in updated) + if unsupported: + raise ValueError( + "P3O behavior sampling does not support additional distribution-truncation parameters: " + f"{', '.join(unsupported)}" + ) updated["temperature"] = _behavior_temperature() return updated @@ -52,3 +72,9 @@ async def generate( behavior_sampling_params(sampling_params, evaluation=evaluation), evaluation=evaluation, ) + + +# _dispatch_generate checks this explicit opt-in before running custom P3O +# generation. The wrapper above rejects known truncation knobs and delegates to +# the built-in SGLang path, which returns sampled-token behavior log-probs. +generate.p3o_behavior_logprob_contract = True diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 3cac95804..6b6057457 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -78,7 +78,7 @@ from ...utils.profile_utils import TrainProfiler from ...utils.training.tensor_backper import TensorBackuper -from .checkpoint import load_checkpoint +from .checkpoint import is_megatron_checkpoint_resume, load_checkpoint from .collective_utils import _agree_drained from .cp_utils import all_gather_with_cp, maybe_padded_total_lengths, slice_with_cp from .data import ( @@ -99,6 +99,7 @@ initial_rollout_policy_snapshot_rollout, maybe_refresh_rollout_policy, rollout_weights_tag, + validate_p3o_periodic_snapshot_resume, validate_update_weights_interval, ) from .weight_update.common import named_params_and_buffers @@ -189,6 +190,7 @@ def _init( # reads model metadata or weights. # Leaf actor with private args, safe to remap in place. prepare_model_maybe_update_args(args) + is_megatron_resume = is_megatron_checkpoint_resume(args.load) and not getattr(args, "finetune", False) self.genrm_manager = None @@ -257,6 +259,11 @@ def _init( # UpdateWeightFromTensor instead of DCS. use_tensor_backuper = not self.args.fully_async or self.args.hybrid update_weights_interval = validate_update_weights_interval(self.args.update_weights_interval) + validate_p3o_periodic_snapshot_resume( + advantage_estimator=getattr(self.args, "advantage_estimator", None), + update_weights_interval=update_weights_interval, + is_megatron_resume=is_megatron_resume, + ) if update_weights_interval > 1 and not use_tensor_backuper: raise ValueError( "update_weights_interval > 1 requires the synchronous or hybrid TensorBackuper weight-update path" @@ -276,7 +283,11 @@ def _init( self.weights_backuper.backup("actor") self._rollout_weights_tag = rollout_weights_tag(update_weights_interval) # Track the rollout at which rollout policy snapshot was created (for observability) - self._rollout_policy_snapshot_rollout = initial_rollout_policy_snapshot_rollout(start_rollout_id) + self._rollout_policy_snapshot_rollout = initial_rollout_policy_snapshot_rollout( + start_rollout_id, + configured_start_rollout_id=self.args.start_rollout_id, + is_megatron_resume=is_megatron_resume, + ) if use_rollout_policy_snapshot: self.weights_backuper.backup(ROLLOUT_POLICY_TAG) diff --git a/relax/backends/megatron/checkpoint.py b/relax/backends/megatron/checkpoint.py index 76752d2f1..f3d38e1b7 100644 --- a/relax/backends/megatron/checkpoint.py +++ b/relax/backends/megatron/checkpoint.py @@ -130,7 +130,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_con exist = Path(load_path).exists() and _is_dir_nonempty(load_path) - if exist and _is_megatron_checkpoint(load_path): + if is_megatron_checkpoint_resume(load_path): _alias_renamed_transfer_queue_enum() try: return _load_checkpoint_megatron( @@ -196,6 +196,14 @@ def _is_megatron_checkpoint(path: str | Path) -> bool: ) +def is_megatron_checkpoint_resume(path: str | Path | None) -> bool: + """Return whether ``path`` is a non-empty native Megatron resume source.""" + if path is None: + return False + candidate = Path(path) + return candidate.is_dir() and _is_dir_nonempty(candidate) and _is_megatron_checkpoint(candidate) + + def _is_hf_checkpoint(path: str | Path) -> bool: return (Path(path) / "config.json").is_file() diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 52effa8ba..cdd3f87c8 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -23,6 +23,7 @@ ) from relax.utils.training.p3o_utils import ( P3OStepContext, + P3OSufficientStats, compute_p3o_sufficient_stats_unchecked, compute_p3o_token_terms, finalize_p3o_step_context, @@ -827,6 +828,8 @@ def get_p3o_context( log_probs: torch.Tensor, behavior_log_probs: torch.Tensor, valid_mask: torch.Tensor, + *, + is_dummy: bool = False, ) -> P3OStepContext: """Resolve the configured P3O ESS scope for one loss micro-batch.""" scope = getattr(args, "p3o_ess_scope", "micro-batch") @@ -835,11 +838,19 @@ def get_p3o_context( if scope != "micro-batch": raise ValueError(f"P3O ESS scope must be 'micro-batch' or 'step', got {scope!r}") - stats, invalid_count = compute_p3o_sufficient_stats_unchecked( - log_probs, - behavior_log_probs, - valid_mask, - ) + if is_dummy: + # Dynamic batching pads shorter DP ranks with dummy micro-batches to + # keep their collective schedule aligned. They still enter the + # collective below, but must add neither ESS moments nor a non-finite + # flag from their placeholder payload. + stats = P3OSufficientStats.zeros(device=log_probs.device) + invalid_count = torch.zeros((), dtype=torch.float64, device=log_probs.device) + else: + stats, invalid_count = compute_p3o_sufficient_stats_unchecked( + log_probs, + behavior_log_probs, + valid_mask, + ) distributed = dist.is_available() and dist.is_initialized() stats = synchronize_p3o_stats( stats, @@ -932,7 +943,19 @@ def p3o_loss_function( dynamic_cp_size=batch.get("dynamic_cp_size", None), dynamic_cp_rank=batch.get("dynamic_cp_rank", None), ) - step_context = get_p3o_context(args, log_probs, behavior_log_probs, valid_mask) + is_dummy = batch.get("__is_dummy__", False) + if is_dummy: + # A dummy batch may carry placeholder behavior log-probs. Exclude every + # token before constructing P3O terms so masking it later with + # ``0.0 * loss`` cannot turn a placeholder NaN into a NaN gradient. + valid_mask = torch.zeros_like(valid_mask, dtype=torch.bool) + step_context = get_p3o_context( + args, + log_probs, + behavior_log_probs, + valid_mask, + is_dummy=is_dummy, + ) terms = compute_p3o_token_terms( log_probs=log_probs, diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index c5bfc26ff..25bd9db32 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -190,6 +190,91 @@ def _preserved_dynamic_cp_group(args: Namespace, model: Sequence[torch.nn.Module inner.pg_collection.cp = original_cp_group +P3O_NO_VALID_TOKENS_ERROR = ( + "P3O: optimizer step has no valid response tokens globally; refusing optimizer and scheduler advancement." +) + + +def _require_p3o_global_valid_tokens( + num_tokens: torch.Tensor | int | None, + device: torch.device, +) -> None: + """Fail a P3O step before gradient normalization can divide by zero. + + The pipeline schedule supplies its exact CP-local token total to + ``finalize_model_grads_func`` after all micro-batches. Reduce it over DP x + CP on the pipeline-last stage, then publish the result to all pipeline + stages so every rank takes the same failure path. The single host read is + intentionally at this optimizer-step boundary rather than in the per-micro- + batch loss hot path. + """ + valid_token_count = torch.zeros((), dtype=torch.float64, device=device) + is_pipeline_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=True) + if is_pipeline_last_stage and num_tokens is not None: + if isinstance(num_tokens, torch.Tensor): + local_num_tokens = num_tokens.detach().to(device=device, dtype=torch.float64) + else: + local_num_tokens = torch.tensor(num_tokens, device=device, dtype=torch.float64) + if local_num_tokens.numel() != 1: + raise ValueError(f"P3O valid-token count must be scalar, got shape {tuple(local_num_tokens.shape)}") + valid_token_count += local_num_tokens.reshape(()) + + distributed = torch.distributed.is_available() and torch.distributed.is_initialized() + if distributed and is_pipeline_last_stage: + torch.distributed.all_reduce( + valid_token_count, + op=torch.distributed.ReduceOp.SUM, + group=mpu.get_data_parallel_group(with_context_parallel=True), + ) + + if distributed and mpu.get_pipeline_model_parallel_world_size() > 1: + pp_group = mpu.get_pipeline_model_parallel_group() + torch.distributed.broadcast( + valid_token_count, + group=pp_group, + group_src=torch.distributed.get_world_size(group=pp_group) - 1, + ) + + if valid_token_count.item() <= 0.0: + raise RuntimeError(P3O_NO_VALID_TOKENS_ERROR) + + +@contextmanager +def _p3o_valid_token_finalizer(args: Namespace, model: Sequence[torch.nn.Module]) -> Iterator[None]: + """Install P3O's zero-token guard immediately before gradient + finalization.""" + if getattr(args, "advantage_estimator", None) != "p3o": + yield + return + + config = get_model_config(model[0]) + original_finalize_model_grads = config.finalize_model_grads_func + if original_finalize_model_grads is None: + raise RuntimeError("P3O requires Megatron's gradient finalizer to be configured before training") + device = next(model[0].parameters()).device + + def finalize_with_p3o_token_guard( + finalizer_model: Sequence[torch.nn.Module], + num_tokens: torch.Tensor | int | None = None, + **kwargs: object, + ) -> None: + _require_p3o_global_valid_tokens(num_tokens, device) + original_finalize_model_grads(finalizer_model, num_tokens, **kwargs) + + config.finalize_model_grads_func = finalize_with_p3o_token_guard + try: + yield + finally: + config.finalize_model_grads_func = original_finalize_model_grads + + +def _require_p3o_step_context_tokens(step_context: object) -> None: + """Reject a globally empty P3O step before its training schedule starts.""" + valid_token_count = getattr(step_context, "valid_token_count") + if valid_token_count.item() <= 0.0: + raise RuntimeError(P3O_NO_VALID_TOKENS_ERROR) + + def _should_use_sft_chunked(args: Namespace) -> bool: """Gate for the SFT chunked-logits path. @@ -1163,9 +1248,10 @@ def forward_step( model=model, num_microbatches=num_microbatches, ) + _require_p3o_step_context_tokens(p3o_step_context) p3o_context_manager = p3o_step_context_published(args, p3o_step_context) - with p3o_context_manager: + with _p3o_valid_token_finalizer(args, model), p3o_context_manager: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=data_iterator, diff --git a/relax/backends/megatron/rollout_policy_lag.py b/relax/backends/megatron/rollout_policy_lag.py index 2c67030e2..f5602f14d 100644 --- a/relax/backends/megatron/rollout_policy_lag.py +++ b/relax/backends/megatron/rollout_policy_lag.py @@ -47,11 +47,49 @@ def compute_rollout_policy_age_rollouts( return current_rollout_id - snapshot_rollout_id -def initial_rollout_policy_snapshot_rollout(start_rollout_id: int) -> int: - """Return the snapshot version aligned with a fresh or resumed run.""" - if start_rollout_id < 0: +def initial_rollout_policy_snapshot_rollout( + backend_start_rollout_id: int, + *, + configured_start_rollout_id: int | None = None, + is_megatron_resume: bool = False, +) -> int: + """Return the snapshot version aligned with the rollout service. + + A cold HuggingFace load reports Megatron iteration zero, so the backend's + next-step value is one while the service correctly starts at its configured + rollout zero. For a cold load, prefer that configured service value; a + native Megatron resume keeps its backend-derived value. + """ + snapshot_rollout_id = ( + backend_start_rollout_id + if is_megatron_resume or configured_start_rollout_id is None + else configured_start_rollout_id + ) + if snapshot_rollout_id < 0: raise ValueError("start_rollout_id must be non-negative") - return start_rollout_id + return snapshot_rollout_id + + +def validate_p3o_periodic_snapshot_resume( + *, + advantage_estimator: str | None, + update_weights_interval: int, + is_megatron_resume: bool, +) -> None: + """Reject P3O resumes that cannot restore the behavior-policy snapshot. + + A periodic rollout-policy snapshot is held only in the in-memory + ``TensorBackuper``. Resuming currently reconstructs it from the current + actor, which changes the behavior policy for the first resumed rollout. + P3O's importance ratio cannot silently absorb that semantic change. + """ + interval = validate_update_weights_interval(update_weights_interval) + if advantage_estimator == "p3o" and interval > 1 and is_megatron_resume: + raise ValueError( + "P3O cannot resume exactly with update_weights_interval > 1 because the periodic rollout-policy " + "snapshot is not checkpointed. Start a fresh run, use update_weights_interval=1, or add " + "snapshot checkpointing." + ) def build_rollout_policy_age_metrics( diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py index 035148aad..73589ee1b 100644 --- a/relax/engine/rollout/sglang_rollout.py +++ b/relax/engine/rollout/sglang_rollout.py @@ -60,6 +60,13 @@ # deadlocking on the same non-reentrant semaphore. A ContextVar is task-local # and propagates down the await chain (including create_task children). _holding_session_lock: contextvars.ContextVar[bool] = contextvars.ContextVar("holding_session_lock", default=False) +_P3O_TRUNCATION_SAMPLING_KEYS = ( + "min_p", + "top_a", + "typical_p", + "epsilon_cutoff", + "eta_cutoff", +) def _ensure_not_holding_session_lock() -> None: @@ -73,6 +80,35 @@ def _ensure_not_holding_session_lock() -> None: ) +def _validate_p3o_behavior_sampling_params( + args: Namespace, + sampling_params: dict[str, Any], + *, + evaluation: bool, +) -> None: + """Reject sampling modes whose returned log-probs are not P3O behavior + policy. + + The built-in generator receives only the CLI-derived parameters, but custom + generation functions can add SGLang's additional distribution-truncation + knobs. Validate the common dispatch boundary before handing parameters to + either path. A custom function additionally declares its contract below. + """ + if evaluation or getattr(args, "advantage_estimator", None) != "p3o": + return + if sampling_params.get("top_p", 1.0) != 1.0 or sampling_params.get("top_k", -1) != -1: + raise ValueError( + "P3O behavior sampling requires top_p=1.0 and top_k=-1 so rollout log-probs describe " + "the untruncated distribution." + ) + unsupported = sorted(key for key in _P3O_TRUNCATION_SAMPLING_KEYS if key in sampling_params) + if unsupported: + raise ValueError( + "P3O behavior sampling does not support additional distribution-truncation parameters: " + f"{', '.join(unsupported)}." + ) + + class GenerateState(metaclass=SingletonMeta): """The global state for the generation process.""" @@ -509,8 +545,19 @@ async def _dispatch_generate( exception would crash the whole rollout step. It is contained here and mapped to an ABORTED sample (mirroring the legacy in-lock abort check). """ + _validate_p3o_behavior_sampling_params(args, sampling_params, evaluation=evaluation) custom_func_path = getattr(sample, "generate_function_path", None) or args.custom_generate_function_path custom_generate_func = load_function(custom_func_path) if custom_func_path is not None else None + if ( + not evaluation + and getattr(args, "advantage_estimator", None) == "p3o" + and custom_generate_func is not None + and not getattr(custom_generate_func, "p3o_behavior_logprob_contract", False) + ): + raise ValueError( + "P3O custom generation must declare p3o_behavior_logprob_contract = True and preserve exact " + "untruncated behavior-policy log-probs for every training token." + ) manages_permit = bool(getattr(custom_generate_func, "manages_inference_permit", False)) async def _run() -> Sample | list[Sample]: diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 1972a1077..701f91a54 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -3,6 +3,7 @@ import argparse import json import os +import re import sys import warnings from typing import Any @@ -75,6 +76,52 @@ def _positive_int(value: str) -> int: return parsed +def _is_megatron_checkpoint_source(path: Any) -> bool: + """Return whether ``path`` is a non-empty native Megatron checkpoint. + + Megatron accepts either a checkpoint root with its latest-iteration + sentinel or a direct ``iter_0000000`` directory. Keep this lightweight + validation-time predicate aligned with the backend loader so a direct + iteration path is never silently downgraded to a HuggingFace cold start. + """ + if path is None or not os.path.isdir(path): + return False + try: + with os.scandir(path) as entries: + if next(entries, None) is None: + return False + except OSError: + return False + return os.path.isfile(os.path.join(path, "latest_checkpointed_iteration.txt")) or bool( + re.fullmatch(r"iter_\d{7}", os.path.basename(os.path.normpath(path))) + ) + + +def _configure_megatron_checkpoint_loading(args: argparse.Namespace) -> None: + """Preserve native checkpoints and configure cold loads consistently.""" + is_megatron_checkpoint = _is_megatron_checkpoint_source(args.load) + if args.megatron_to_hf_mode == "bridge": + if is_megatron_checkpoint: + # Native checkpoints are loaded directly by Megatron rather than + # through the HuggingFace bridge. + return + if args.load is None: + args.load = args.ref_load or args.hf_checkpoint + # A HuggingFace/cold load starts the rollout service at zero. + args.start_rollout_id = 0 + return + + if is_megatron_checkpoint: + return + args.no_load_optim = True + args.no_load_rng = True + args.finetune = True + args.load = args.ref_load + if args.ref_ckpt_step is not None: + args.ckpt_step = args.ref_ckpt_step + args.start_rollout_id = 0 + + def reset_arg(parser, name, **kwargs): """Reset the default value of a Megatron argument. @@ -1839,7 +1886,10 @@ def add_algo_arguments(parser): "--p3o-ess-scope", choices=["micro-batch", "step"], default="micro-batch", - help="P3O ESS scope: paper-compatible micro-batch (default) or optimizer step.", + help=( + "P3O ESS scope: paper-compatible but topology-dependent micro-batch (default), " + "or optimizer step for a partition-invariant adaptive cap." + ), ) parser.add_argument( "--p3o-kl-mode", @@ -2930,6 +2980,14 @@ def _validate_p3o_args(args: argparse.Namespace) -> None: "reintroduces a per-micro-batch denominator, so the loss would depend on " "how the optimizer step is split into micro-batches." ) + rollout_top_p = getattr(args, "rollout_top_p", 1.0) + rollout_top_k = getattr(args, "rollout_top_k", -1) + if rollout_top_p != 1.0 or rollout_top_k != -1: + raise ValueError( + "P3O requires behavior log-probs from the untruncated rollout distribution. " + "Use --rollout-top-p 1.0 and --rollout-top-k -1; truncated sampling changes the behavior policy " + "unless its exact normalized log-probs are explicitly supported." + ) if args.use_tis: raise ValueError( "P3O and TIS (--use-tis) are mutually exclusive: both correct the same " @@ -2955,6 +3013,12 @@ def _validate_p3o_args(args: argparse.Namespace) -> None: if getattr(args, "custom_pg_loss_reducer_function_path", None) is not None: raise ValueError("P3O requires token-sum normalization and does not support a custom PG-loss reducer.") + if scope == "micro-batch" and getattr(args, "recompute_loss_function", False): + raise ValueError( + "P3O micro-batch ESS reduces statistics inside the loss callback, which checkpoint replay would execute " + "again during backward. Disable --recompute-loss-function or use --p3o-ess-scope step." + ) + if scope == "step": # The ESS pre-pass replays the same micro-batch window under no_grad. Ops # that mutate state on a forward would make the two passes disagree. @@ -3178,32 +3242,7 @@ def slime_validate_args(args): validate_opd_args(args, is_sft=is_sft, log=logger) - if args.megatron_to_hf_mode == "bridge": - if ( - args.load is not None - and os.path.exists(args.load) - and os.path.exists(os.path.join(args.load, "latest_checkpointed_iteration.txt")) - ): - # If is a Megatron checkpoint, won't use bridge to load hf weight. - pass - else: - if args.load is None: - args.load = args.ref_load or args.hf_checkpoint - # If is a HF checkpoint, set start_rollout_id to 0 here. - args.start_rollout_id = 0 - else: - if ( - args.load is None - or not os.path.exists(args.load) - or not os.path.exists(os.path.join(args.load, "latest_checkpointed_iteration.txt")) - ): - args.no_load_optim = True - args.no_load_rng = True - args.finetune = True - args.load = args.ref_load - if args.ref_ckpt_step is not None: - args.ckpt_step = args.ref_ckpt_step - args.start_rollout_id = 0 + _configure_megatron_checkpoint_loading(args) if args.eval_interval is not None: if args.loss_type == "sft": diff --git a/tests/backends/megatron/test_p3o_distributed.py b/tests/backends/megatron/test_p3o_distributed.py index 60df65d25..b86f245b7 100644 --- a/tests/backends/megatron/test_p3o_distributed.py +++ b/tests/backends/megatron/test_p3o_distributed.py @@ -7,24 +7,34 @@ import math import os import socket +import sys +from types import ModuleType +from unittest.mock import patch import torch import torch.distributed as dist import torch.multiprocessing as mp -from tests.backends.megatron._megatron_stub import stubbed_megatron_modules - - -with stubbed_megatron_modules(("megatron", "ray", "tensordict")): - from relax.backends.megatron import p3o_step - from relax.backends.megatron.p3o_step import synchronize_p3o_stats - from relax.utils.training.p3o_utils import ( P3OSufficientStats, compute_p3o_sufficient_stats, compute_p3o_token_terms, finalize_p3o_step_context, ) +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules + + +stream_dataloader = ModuleType("relax.utils.data.stream_dataloader") +stream_dataloader.StreamingTQIterator = object + +with ( + patch.dict(sys.modules, {"relax.utils.data.stream_dataloader": stream_dataloader}), + stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), +): + from relax.backends.megatron import loss as loss_module + from relax.backends.megatron import model as model_module + from relax.backends.megatron import p3o_step + from relax.backends.megatron.p3o_step import synchronize_p3o_stats def _free_port() -> int: @@ -197,6 +207,115 @@ def assert_partition(shards: list[torch.Tensor], process_group) -> None: dist.destroy_process_group() +def _zero_token_finalizer_worker(rank: int, world_size: int, port: int) -> None: + _init_gloo(rank, world_size, port) + try: + pp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + dp_first_group = dist.new_group([0, 2]) + dp_last_group = dist.new_group([1, 3]) + is_last_stage = rank in {1, 3} + model_module.mpu.is_pipeline_last_stage = lambda ignore_virtual=True: is_last_stage + model_module.mpu.get_pipeline_model_parallel_world_size = lambda: 2 + model_module.mpu.get_pipeline_model_parallel_group = lambda: pp_groups[rank // 2] + model_module.mpu.get_data_parallel_group = lambda with_context_parallel=True: ( + dp_last_group if is_last_stage else dp_first_group + ) + + def assert_outcome(num_tokens: float, *, should_raise: bool) -> None: + raised = False + try: + model_module._require_p3o_global_valid_tokens( + torch.tensor(num_tokens), + torch.device("cpu"), + ) + except RuntimeError as error: + assert "no valid response tokens globally" in str(error) + raised = True + + outcomes = torch.tensor(float(raised)) + dist.all_reduce(outcomes) + assert outcomes.item() == (world_size if should_raise else 0) + + healthy = torch.ones(()) + dist.all_reduce(healthy) + assert healthy.item() == world_size + + # Both DP shards have no tokens: every PP stage must fail before the + # Megatron finalizer can normalize gradients by zero. + assert_outcome(0.0, should_raise=True) + # Only one DP shard contributes valid tokens; the DP reduction must + # publish a positive count to every PP stage. + assert_outcome(2.0 if rank == 3 else 0.0, should_raise=False) + finally: + dist.destroy_process_group() + + +def _real_dummy_dp_worker(rank: int, world_size: int, port: int) -> None: + _init_gloo(rank, world_size, port) + try: + loss_module.mpu.is_pipeline_last_stage = lambda ignore_virtual=True: True + loss_module.mpu.get_data_parallel_group = lambda with_context_parallel=True: dist.group.WORLD + loss_module.mpu.get_pipeline_model_parallel_world_size = lambda: 1 + + behavior = torch.full((4,), -2.0) + ratios = torch.tensor([1.0, 2.0, 0.5, 4.0]) + real_log_probs = behavior + ratios.log() + advantages = torch.tensor([1.0, -1.0, 0.5, 2.0]) + real_mask = torch.tensor([True, True, False, True]) + args = type("Args", (), {"p3o_ess_scope": "micro-batch"})() + + oracle_context = finalize_p3o_step_context(compute_p3o_sufficient_stats(real_log_probs, behavior, real_mask)) + oracle_log_probs = real_log_probs.clone().requires_grad_(True) + oracle_terms = compute_p3o_token_terms( + oracle_log_probs, + behavior, + advantages, + real_mask, + oracle_context, + ) + oracle_loss = (oracle_terms.score_loss + oracle_terms.adaptive_kl_loss).sum() + oracle_loss = oracle_loss / oracle_context.valid_token_count + oracle_loss.backward() + + is_dummy = rank == 1 + local_log_probs = real_log_probs.clone().requires_grad_(True) + local_behavior = torch.full_like(behavior, float("nan")) if is_dummy else behavior + local_mask = torch.zeros_like(real_mask) if is_dummy else real_mask + context = loss_module.get_p3o_context( + args, + local_log_probs, + local_behavior, + real_mask, + is_dummy=is_dummy, + ) + torch.testing.assert_close(context.adaptive_cap, oracle_context.adaptive_cap) + + local_terms = compute_p3o_token_terms( + local_log_probs, + local_behavior, + advantages, + local_mask, + context, + ) + local_loss = (local_terms.score_loss + local_terms.adaptive_kl_loss).sum() + local_loss = local_loss / context.valid_token_count + 0.0 * local_log_probs.sum() + assert torch.isfinite(local_loss) + local_loss.backward() + + reduced_loss = local_loss.detach().clone() + reduced_gradient = local_log_probs.grad.detach().clone() + dist.all_reduce(reduced_loss) + dist.all_reduce(reduced_gradient) + torch.testing.assert_close(reduced_loss, oracle_loss.detach()) + torch.testing.assert_close(reduced_gradient, oracle_log_probs.grad.detach()) + + healthy = torch.ones(()) + dist.all_reduce(healthy) + assert healthy.item() == world_size + finally: + dist.destroy_process_group() + + def test_p3o_distributed_nonfinite_fails_synchronously(): world_size = 2 mp.spawn(_nonfinite_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) @@ -210,3 +329,13 @@ def test_p3o_distributed_pipeline_broadcasts_last_stage_stats(): def test_p3o_distributed_partition_and_objective_invariance(): world_size = 4 mp.spawn(_partition_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) + + +def test_p3o_distributed_finalizer_guard_is_pp_and_dp_symmetric(): + world_size = 4 + mp.spawn(_zero_token_finalizer_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) + + +def test_p3o_distributed_real_and_dummy_dp_match_real_only_oracle(): + world_size = 2 + mp.spawn(_real_dummy_dp_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) diff --git a/tests/backends/megatron/test_p3o_loss.py b/tests/backends/megatron/test_p3o_loss.py index ab0ad69f9..a28649df7 100644 --- a/tests/backends/megatron/test_p3o_loss.py +++ b/tests/backends/megatron/test_p3o_loss.py @@ -16,7 +16,7 @@ from tests.backends.megatron._megatron_stub import stubbed_megatron_modules -with stubbed_megatron_modules(("megatron", "ray")): +with stubbed_megatron_modules(("megatron", "ray", "tensordict")): from relax.backends.megatron import loss as loss_module from relax.utils.training.p3o_utils import P3OStepContext @@ -52,6 +52,28 @@ def test_get_p3o_context_computes_micro_batch_scope_without_prepass(): assert torch.equal(context.normalized_ess, context.adaptive_cap) +def test_get_p3o_context_zeroes_dummy_micro_batch_before_collective(monkeypatch): + captured = {} + + def capture_stats(stats, invalid_count, **kwargs): + captured["stats"] = stats.as_vector() + captured["invalid_count"] = invalid_count + return stats + + monkeypatch.setattr(loss_module, "synchronize_p3o_stats", capture_stats) + context = loss_module.get_p3o_context( + Namespace(p3o_ess_scope="micro-batch"), + torch.tensor([float("nan")]), + torch.tensor([0.0]), + torch.tensor([True]), + is_dummy=True, + ) + + assert torch.equal(captured["stats"], torch.zeros(3, dtype=torch.float64)) + assert torch.equal(captured["invalid_count"], torch.zeros((), dtype=torch.float64)) + assert context.valid_token_count.item() == 0 + + def test_get_p3o_context_rejects_unknown_scope(): args = Namespace(p3o_ess_scope="window") @@ -91,21 +113,34 @@ def test_p3o_loss_reports_complete_schema_without_reference_kl(monkeypatch): "get_cp_local_valid_mask", lambda *args, **kwargs: torch.tensor([True, True]), ) + dummy_masks = [] + dummy_flags = [] + + def capture_context(*args, is_dummy=False): + dummy_masks.append(args[3]) + dummy_flags.append(is_dummy) + return step_context + + monkeypatch.setattr(loss_module, "get_p3o_context", capture_context) batch = { "advantages": torch.tensor([1.0, -1.0]), - "rollout_log_probs": [log_probs.detach().clone()], + "rollout_log_probs": [torch.full_like(log_probs.detach(), float("nan"))], "unconcat_tokens": [torch.tensor([1, 2])], "total_lengths": [2], "response_lengths": [2], "loss_masks": [torch.ones(2)], + "__is_dummy__": True, } - _, metrics = loss_module.p3o_loss_function(args, batch, torch.zeros(1), torch.sum) + loss, metrics = loss_module.p3o_loss_function(args, batch, torch.zeros(1), torch.sum) assert REQUIRED_P3O_METRICS <= metrics.keys() assert not any(metric.startswith("opd/") for metric in metrics) assert torch.equal(metrics["p3o/reference_kl"], torch.zeros(())) assert not metrics["p3o/reference_kl"].requires_grad + assert torch.isfinite(loss) + assert torch.equal(dummy_masks[0], torch.zeros(2, dtype=torch.bool)) + assert dummy_flags == [True] def test_p3o_loss_function_normalizes_by_true_valid_tokens(monkeypatch): diff --git a/tests/backends/megatron/test_p3o_model_step.py b/tests/backends/megatron/test_p3o_model_step.py index 4499cbfea..88dbd436c 100644 --- a/tests/backends/megatron/test_p3o_model_step.py +++ b/tests/backends/megatron/test_p3o_model_step.py @@ -12,6 +12,7 @@ from unittest.mock import patch import pytest +import torch from tests.backends.megatron._megatron_stub import stubbed_megatron_modules @@ -25,7 +26,10 @@ patch.dict(sys.modules, {"relax.utils.data.stream_dataloader": stream_dataloader}), stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), ): - from relax.backends.megatron.model import _preserved_dynamic_cp_group + from relax.backends.megatron import model as model_module + + +_preserved_dynamic_cp_group = model_module._preserved_dynamic_cp_group def test_p3o_model_step_restores_dynamic_cp_group_after_error(): @@ -43,6 +47,71 @@ def test_p3o_model_step_restores_dynamic_cp_group_after_error(): assert inner.pg_collection.cp is original_group +def test_p3o_model_step_rejects_global_zero_valid_tokens(monkeypatch): + monkeypatch.setattr(model_module.mpu, "is_pipeline_last_stage", lambda **kwargs: True) + monkeypatch.setattr(model_module.torch.distributed, "is_available", lambda: False) + + with pytest.raises(RuntimeError, match="no valid response tokens globally"): + model_module._require_p3o_global_valid_tokens( + torch.tensor(0.0), + torch.device("cpu"), + ) + + +def test_p3o_model_step_allows_positive_global_valid_tokens(monkeypatch): + monkeypatch.setattr(model_module.mpu, "is_pipeline_last_stage", lambda **kwargs: True) + monkeypatch.setattr(model_module.torch.distributed, "is_available", lambda: False) + + model_module._require_p3o_global_valid_tokens( + torch.tensor(2.0), + torch.device("cpu"), + ) + + +def test_p3o_model_step_rejects_zero_step_context_before_training_schedule(): + with pytest.raises(RuntimeError, match="no valid response tokens globally"): + model_module._require_p3o_step_context_tokens(SimpleNamespace(valid_token_count=torch.tensor(0.0))) + + model_module._require_p3o_step_context_tokens(SimpleNamespace(valid_token_count=torch.tensor(1.0))) + + +def test_p3o_model_step_finalizer_guard_skips_original_finalizer_on_zero_tokens(monkeypatch): + original_calls = [] + config = SimpleNamespace( + finalize_model_grads_func=lambda *args, **kwargs: original_calls.append((args, kwargs)), + ) + model = [torch.nn.Linear(1, 1)] + monkeypatch.setattr(model_module, "get_model_config", lambda _: config) + monkeypatch.setattr(model_module.mpu, "is_pipeline_last_stage", lambda **kwargs: True) + monkeypatch.setattr(model_module.torch.distributed, "is_available", lambda: False) + original_finalizer = config.finalize_model_grads_func + + with pytest.raises(RuntimeError, match="no valid response tokens globally"): + with model_module._p3o_valid_token_finalizer(Namespace(advantage_estimator="p3o"), model): + config.finalize_model_grads_func(model, torch.tensor(0.0)) + + assert original_calls == [] + assert config.finalize_model_grads_func is original_finalizer + + +def test_p3o_model_step_finalizer_guard_calls_original_and_restores_config(monkeypatch): + original_calls = [] + config = SimpleNamespace( + finalize_model_grads_func=lambda *args, **kwargs: original_calls.append((args, kwargs)), + ) + model = [torch.nn.Linear(1, 1)] + monkeypatch.setattr(model_module, "get_model_config", lambda _: config) + monkeypatch.setattr(model_module.mpu, "is_pipeline_last_stage", lambda **kwargs: True) + monkeypatch.setattr(model_module.torch.distributed, "is_available", lambda: False) + original_finalizer = config.finalize_model_grads_func + + with model_module._p3o_valid_token_finalizer(Namespace(advantage_estimator="p3o"), model): + config.finalize_model_grads_func(model, torch.tensor(2.0), pg_collection="pg") + + assert original_calls == [((model, torch.tensor(2.0)), {"pg_collection": "pg"})] + assert config.finalize_model_grads_func is original_finalizer + + def test_p3o_model_step_guard_covers_stats_and_train_passes(): tree = ast.parse(MODEL_PATH.read_text(encoding="utf-8")) train_one_step = next( @@ -67,3 +136,27 @@ def test_p3o_model_step_guard_covers_stats_and_train_passes(): assert "p3o_ess_scope" in guarded_source assert "micro-batch" in guarded_source assert "step" in guarded_source + + +def test_p3o_model_step_wraps_gradient_finalization_with_valid_token_guard(): + tree = ast.parse(MODEL_PATH.read_text(encoding="utf-8")) + train_one_step = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "train_one_step" + ) + p3o_finalizer_guard = next( + node + for node in ast.walk(train_one_step) + if isinstance(node, ast.With) + and any( + isinstance(child, ast.Name) and child.id == "_p3o_valid_token_finalizer" + for item in node.items + for child in ast.walk(item.context_expr) + ) + ) + guarded_calls = { + child.func.id + for child in ast.walk(p3o_finalizer_guard) + if isinstance(child, ast.Call) and isinstance(child.func, ast.Name) + } + + assert "forward_backward_func" in guarded_calls diff --git a/tests/backends/megatron/test_p3o_observability.py b/tests/backends/megatron/test_p3o_observability.py index 0c9d50304..5c4496190 100644 --- a/tests/backends/megatron/test_p3o_observability.py +++ b/tests/backends/megatron/test_p3o_observability.py @@ -3,6 +3,9 @@ """Behavior tests for P3O rollout-policy age observability.""" import ast +import os +import re +import types from pathlib import Path import pytest @@ -15,10 +18,31 @@ maybe_refresh_rollout_policy, rollout_weights_tag, should_refresh_rollout_policy, + validate_p3o_periodic_snapshot_resume, validate_update_weights_interval, ) +CHECKPOINT_PATH = Path(__file__).resolve().parents[3] / "relax" / "backends" / "megatron" / "checkpoint.py" + + +def _load_megatron_resume_detector(): + """Extract the path-only resume detector without importing Megatron.""" + tree = ast.parse(CHECKPOINT_PATH.read_text(encoding="utf-8")) + names = {"_is_dir_nonempty", "_is_megatron_checkpoint", "is_megatron_checkpoint_resume"} + funcs = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + module = types.ModuleType("_megatron_resume_detector") + module.Path = Path + module.os = os + module.re = re + exec(compile(ast.Module(body=funcs, type_ignores=[]), str(CHECKPOINT_PATH), "exec"), module.__dict__) + return module + + +megatron_resume_detector = _load_megatron_resume_detector() +is_megatron_checkpoint_resume = megatron_resume_detector.is_megatron_checkpoint_resume + + class _RecordingBackuper: def __init__(self) -> None: self.copies: list[tuple[str, str]] = [] @@ -74,11 +98,116 @@ def test_rollout_policy_snapshot_initializes_for_fresh_and_resumed_runs(): assert compute_rollout_policy_age_rollouts(101, initial_rollout_policy_snapshot_rollout(101)) == 0 +def test_rollout_policy_snapshot_uses_service_start_for_cold_hf_load(): + snapshot_rollout = initial_rollout_policy_snapshot_rollout( + 1, + configured_start_rollout_id=0, + is_megatron_resume=False, + ) + + assert snapshot_rollout == 0 + assert compute_rollout_policy_age_rollouts(0, snapshot_rollout) == 0 + + def test_rollout_policy_snapshot_rejects_invalid_resume_version(): with pytest.raises(ValueError, match="start_rollout_id"): initial_rollout_policy_snapshot_rollout(-1) +@pytest.mark.parametrize( + ("advantage_estimator", "update_weights_interval", "is_megatron_resume"), + [ + ("p3o", 1, True), + ("p3o", 3, False), + ("grpo", 3, True), + ], +) +def test_p3o_periodic_snapshot_resume_allows_exactly_reconstructible_cases( + advantage_estimator, + update_weights_interval, + is_megatron_resume, +): + validate_p3o_periodic_snapshot_resume( + advantage_estimator=advantage_estimator, + update_weights_interval=update_weights_interval, + is_megatron_resume=is_megatron_resume, + ) + + +def test_p3o_periodic_snapshot_resume_rejects_inexact_resume(): + with pytest.raises(ValueError, match="cannot resume exactly"): + validate_p3o_periodic_snapshot_resume( + advantage_estimator="p3o", + update_weights_interval=3, + is_megatron_resume=True, + ) + + +def test_native_megatron_resume_detection_distinguishes_hf_and_fresh_paths(tmp_path): + checkpoint_root = tmp_path / "checkpoint" + checkpoint_root.mkdir() + (checkpoint_root / "latest_checkpointed_iteration.txt").write_text("12", encoding="utf-8") + assert is_megatron_checkpoint_resume(checkpoint_root) + + iteration_checkpoint = tmp_path / "iter_0000000" + iteration_checkpoint.mkdir() + (iteration_checkpoint / "common.pt").write_bytes(b"checkpoint") + assert is_megatron_checkpoint_resume(iteration_checkpoint) + + hf_checkpoint = tmp_path / "hf" + hf_checkpoint.mkdir() + (hf_checkpoint / "config.json").write_text("{}", encoding="utf-8") + assert not is_megatron_checkpoint_resume(hf_checkpoint) + assert not is_megatron_checkpoint_resume(tmp_path / "empty") + assert not is_megatron_checkpoint_resume(tmp_path / "missing") + assert not is_megatron_checkpoint_resume(None) + + +def test_actor_derives_periodic_resume_guard_from_native_checkpoint_path(): + actor_path = Path(__file__).resolve().parents[3] / "relax" / "backends" / "megatron" / "actor.py" + tree = ast.parse(actor_path.read_text(encoding="utf-8")) + actor_class = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "MegatronTrainRayActor" + ) + init_method = next(node for node in actor_class.body if isinstance(node, ast.FunctionDef) and node.name == "_init") + + native_resume_assignment = next( + node + for node in ast.walk(init_method) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "is_megatron_resume" for target in node.targets) + ) + native_resume_source = ast.unparse(native_resume_assignment.value) + assert "is_megatron_checkpoint_resume" in native_resume_source + assert "finetune" in native_resume_source + + resume_validator = next( + node + for node in ast.walk(init_method) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "validate_p3o_periodic_snapshot_resume" + ) + is_megatron_resume_keyword = next( + keyword for keyword in resume_validator.keywords if keyword.arg == "is_megatron_resume" + ) + assert isinstance(is_megatron_resume_keyword.value, ast.Name) + assert is_megatron_resume_keyword.value.id == "is_megatron_resume" + + snapshot_initializer = next( + node + for node in ast.walk(init_method) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "initial_rollout_policy_snapshot_rollout" + ) + snapshot_keywords = {keyword.arg: keyword.value for keyword in snapshot_initializer.keywords} + assert isinstance(snapshot_keywords["configured_start_rollout_id"], ast.Attribute) + assert snapshot_keywords["configured_start_rollout_id"].attr == "start_rollout_id" + assert isinstance(snapshot_keywords["is_megatron_resume"], ast.Name) + assert snapshot_keywords["is_megatron_resume"].id == "is_megatron_resume" + + def test_rollout_policy_age_metrics_have_exact_keys_and_values(): assert build_rollout_policy_age_metrics(current_rollout_id=7, rollout_policy_snapshot_rollout=5) == { "train/current_rollout_id": 7, diff --git a/tests/engine/rollout/test_p3o_sampling_contract.py b/tests/engine/rollout/test_p3o_sampling_contract.py new file mode 100644 index 000000000..290ea7c40 --- /dev/null +++ b/tests/engine/rollout/test_p3o_sampling_contract.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for P3O's rollout behavior-policy sampling contract.""" + +import ast +import types +from argparse import Namespace +from pathlib import Path +from typing import Any + +import pytest + + +SGLANG_ROLLOUT_PATH = Path(__file__).resolve().parents[3] / "relax" / "engine" / "rollout" / "sglang_rollout.py" + + +def _load_p3o_sampling_validator(): + """Extract the dependency-free validation helper from the rollout source.""" + tree = ast.parse(SGLANG_ROLLOUT_PATH.read_text(encoding="utf-8")) + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "_P3O_TRUNCATION_SAMPLING_KEYS" for target in node.targets) + ) + validator = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_validate_p3o_behavior_sampling_params" + ) + module = types.ModuleType("_p3o_sampling_contract") + module.Any = Any + module.Namespace = Namespace + exec(compile(ast.Module(body=[assignment, validator], type_ignores=[]), str(SGLANG_ROLLOUT_PATH), "exec"), module.__dict__) + return module._validate_p3o_behavior_sampling_params + + +validate_p3o_behavior_sampling_params = _load_p3o_sampling_validator() + + +def test_p3o_sampling_contract_allows_only_untruncated_training_sampling(): + args = Namespace(advantage_estimator="p3o") + + validate_p3o_behavior_sampling_params(args, {"top_p": 1.0, "top_k": -1}, evaluation=False) + + for sampling_params in ({"top_p": 0.9}, {"top_k": 32}, {"min_p": 0.1}): + with pytest.raises(ValueError, match="P3O behavior sampling"): + validate_p3o_behavior_sampling_params(args, sampling_params, evaluation=False) + + +def test_p3o_sampling_contract_leaves_evaluation_and_non_p3o_unchanged(): + validate_p3o_behavior_sampling_params( + Namespace(advantage_estimator="p3o"), + {"top_p": 0.9, "min_p": 0.1}, + evaluation=True, + ) + validate_p3o_behavior_sampling_params( + Namespace(advantage_estimator="grpo"), + {"top_p": 0.9, "min_p": 0.1}, + evaluation=False, + ) + + +def test_p3o_dispatch_requires_an_explicit_custom_generation_contract(): + tree = ast.parse(SGLANG_ROLLOUT_PATH.read_text(encoding="utf-8")) + dispatch = next(node for node in tree.body if isinstance(node, ast.AsyncFunctionDef) and node.name == "_dispatch_generate") + string_constants = { + node.value + for node in ast.walk(dispatch) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + + assert "p3o_behavior_logprob_contract" in string_constants diff --git a/tests/examples/algorithms/p3o/test_rollout.py b/tests/examples/algorithms/p3o/test_rollout.py index 02e1abba6..7cedbf659 100644 --- a/tests/examples/algorithms/p3o/test_rollout.py +++ b/tests/examples/algorithms/p3o/test_rollout.py @@ -9,14 +9,26 @@ from examples.algorithms.p3o import rollout -def test_behavior_sampling_params_overrides_only_temperature(monkeypatch): +def test_behavior_sampling_params_overrides_temperature_for_untruncated_sampling(monkeypatch): monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", "0.6") - original = {"temperature": 1.0, "top_p": 0.9, "max_new_tokens": 64} + original = {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 64} updated = rollout.behavior_sampling_params(original, evaluation=False) - assert updated == {"temperature": 0.6, "top_p": 0.9, "max_new_tokens": 64} - assert original == {"temperature": 1.0, "top_p": 0.9, "max_new_tokens": 64} + assert updated == {"temperature": 0.6, "top_p": 1.0, "top_k": -1, "max_new_tokens": 64} + assert original == {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 64} + + +@pytest.mark.parametrize("sampling_params", [{"top_p": 0.9}, {"top_k": 32}, {"min_p": 0.1}]) +def test_behavior_sampling_params_rejects_truncated_training_sampling(monkeypatch, sampling_params): + monkeypatch.setenv("P3O_BEHAVIOR_TEMPERATURE", "0.6") + + with pytest.raises(ValueError, match="P3O behavior sampling"): + rollout.behavior_sampling_params(sampling_params, evaluation=False) + + +def test_p3o_rollout_declares_behavior_logprob_contract(): + assert rollout.generate.p3o_behavior_logprob_contract is True @pytest.mark.parametrize(("raw_value", "expected"), [("0.6", 0.6), ("1.2", 1.2), ("2.0", 2.0)]) @@ -77,7 +89,7 @@ async def fake_generate(args, sample, sampling_params, evaluation=False): monkeypatch.setattr(rollout, "_sglang_generate", fake_generate) args = SimpleNamespace() sample = object() - original = {"temperature": 1.0, "top_p": 0.95, "max_new_tokens": 32} + original = {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 32} result = await rollout.generate(args, sample, original, evaluation=False) @@ -85,7 +97,7 @@ async def fake_generate(args, sample, sampling_params, evaluation=False): assert captured == { "args": args, "sample": sample, - "sampling_params": {"temperature": 1.2, "top_p": 0.95, "max_new_tokens": 32}, + "sampling_params": {"temperature": 1.2, "top_p": 1.0, "top_k": -1, "max_new_tokens": 32}, "evaluation": False, } - assert original == {"temperature": 1.0, "top_p": 0.95, "max_new_tokens": 32} + assert original == {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 32} diff --git a/tests/utils/test_p3o_arguments.py b/tests/utils/test_p3o_arguments.py index 230836a29..d94bd250f 100644 --- a/tests/utils/test_p3o_arguments.py +++ b/tests/utils/test_p3o_arguments.py @@ -13,10 +13,14 @@ module source by AST rather than imported. """ +import argparse import ast +import os +import re import types from argparse import Namespace from pathlib import Path +from typing import Any import pytest @@ -26,8 +30,6 @@ def _load_validator(): """Extract ``_validate_p3o_args`` without importing arguments.py.""" - import argparse - tree = ast.parse(ARGUMENTS_PATH.read_text(encoding="utf-8")) func = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_validate_p3o_args") module = types.ModuleType("_p3o_args") @@ -39,6 +41,24 @@ def _load_validator(): validate_p3o_args = _load_validator() +def _load_checkpoint_loading_helpers(): + """Extract checkpoint-loading helpers without importing optional + backends.""" + tree = ast.parse(ARGUMENTS_PATH.read_text(encoding="utf-8")) + names = {"_is_megatron_checkpoint_source", "_configure_megatron_checkpoint_loading"} + funcs = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + module = types.ModuleType("_checkpoint_loading_args") + module.Any = Any + module.argparse = argparse + module.os = os + module.re = re + exec(compile(ast.Module(body=funcs, type_ignores=[]), str(ARGUMENTS_PATH), "exec"), module.__dict__) + return module + + +checkpoint_loading_helpers = _load_checkpoint_loading_helpers() + + def _p3o_kl_mode_choices() -> list[str]: """Read the CLI choices without importing Relax's Megatron dependency chain.""" @@ -67,6 +87,9 @@ def _p3o_args(**overrides) -> Namespace: clip_high=0.2, use_rollout_logprobs=True, calculate_per_token_loss=True, + rollout_top_p=1.0, + rollout_top_k=-1, + recompute_loss_function=False, use_tis=False, true_on_policy_mode=False, use_critic=False, @@ -138,6 +161,76 @@ def test_p3o_arguments_step_scope_accepts_inactive_lora_dropout(): validate_p3o_args(_p3o_args(p3o_ess_scope="step", lora_rank=0, lora_dropout=0.1)) +@pytest.mark.parametrize( + "overrides", + [ + dict(rollout_top_p=0.9), + dict(rollout_top_k=32), + ], +) +def test_p3o_arguments_rejects_truncated_rollout_behavior_logprobs(overrides): + with pytest.raises(ValueError, match="untruncated rollout distribution"): + validate_p3o_args(_p3o_args(**overrides)) + + +def test_p3o_arguments_rejects_micro_batch_ess_with_recomputed_loss(): + with pytest.raises(ValueError, match="checkpoint replay"): + validate_p3o_args(_p3o_args(recompute_loss_function=True)) + + +def test_p3o_arguments_step_scope_allows_recomputed_loss(): + validate_p3o_args(_p3o_args(p3o_ess_scope="step", recompute_loss_function=True)) + + +def test_megatron_checkpoint_source_recognizes_root_and_direct_iteration_paths(tmp_path): + checkpoint_root = tmp_path / "checkpoint" + checkpoint_root.mkdir() + (checkpoint_root / "latest_checkpointed_iteration.txt").write_text("12", encoding="utf-8") + assert checkpoint_loading_helpers._is_megatron_checkpoint_source(checkpoint_root) + + iteration_checkpoint = tmp_path / "iter_0000000" + iteration_checkpoint.mkdir() + (iteration_checkpoint / "common.pt").write_bytes(b"checkpoint") + assert checkpoint_loading_helpers._is_megatron_checkpoint_source(iteration_checkpoint) + + hf_checkpoint = tmp_path / "hf" + hf_checkpoint.mkdir() + (hf_checkpoint / "config.json").write_text("{}", encoding="utf-8") + empty_checkpoint = tmp_path / "iter_0000001" + empty_checkpoint.mkdir() + assert not checkpoint_loading_helpers._is_megatron_checkpoint_source(hf_checkpoint) + assert not checkpoint_loading_helpers._is_megatron_checkpoint_source(empty_checkpoint) + assert not checkpoint_loading_helpers._is_megatron_checkpoint_source(tmp_path / "missing") + + +@pytest.mark.parametrize("megatron_to_hf_mode", ["bridge", "raw"]) +def test_megatron_checkpoint_loading_preserves_direct_native_iteration_resume(tmp_path, megatron_to_hf_mode): + iteration_checkpoint = tmp_path / "iter_0000000" + iteration_checkpoint.mkdir() + (iteration_checkpoint / "common.pt").write_bytes(b"checkpoint") + args = Namespace( + megatron_to_hf_mode=megatron_to_hf_mode, + load=str(iteration_checkpoint), + ref_load="ref-checkpoint", + hf_checkpoint="hf-checkpoint", + start_rollout_id=None, + no_load_optim=False, + no_load_rng=False, + finetune=False, + ref_ckpt_step=9, + ckpt_step=None, + ) + + checkpoint_loading_helpers._configure_megatron_checkpoint_loading(args) + + assert args.load == str(iteration_checkpoint) + assert args.start_rollout_id is None + assert not args.no_load_optim + assert not args.no_load_rng + assert not args.finetune + assert args.ckpt_step is None + + @pytest.mark.parametrize( ("reason", "overrides"), [ From 8679b342f097548c5e40f086ae025f1f65e1f59a Mon Sep 17 00:00:00 2001 From: Chream Date: Fri, 14 Aug 2026 10:53:30 +0800 Subject: [PATCH 05/30] test: format P3O sampling contract test --- .../rollout/test_p3o_sampling_contract.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/engine/rollout/test_p3o_sampling_contract.py b/tests/engine/rollout/test_p3o_sampling_contract.py index 290ea7c40..480d8bd0a 100644 --- a/tests/engine/rollout/test_p3o_sampling_contract.py +++ b/tests/engine/rollout/test_p3o_sampling_contract.py @@ -15,13 +15,16 @@ def _load_p3o_sampling_validator(): - """Extract the dependency-free validation helper from the rollout source.""" + """Extract the dependency-free validation helper from the rollout + source.""" tree = ast.parse(SGLANG_ROLLOUT_PATH.read_text(encoding="utf-8")) assignment = next( node for node in tree.body if isinstance(node, ast.Assign) - and any(isinstance(target, ast.Name) and target.id == "_P3O_TRUNCATION_SAMPLING_KEYS" for target in node.targets) + and any( + isinstance(target, ast.Name) and target.id == "_P3O_TRUNCATION_SAMPLING_KEYS" for target in node.targets + ) ) validator = next( node @@ -31,7 +34,10 @@ def _load_p3o_sampling_validator(): module = types.ModuleType("_p3o_sampling_contract") module.Any = Any module.Namespace = Namespace - exec(compile(ast.Module(body=[assignment, validator], type_ignores=[]), str(SGLANG_ROLLOUT_PATH), "exec"), module.__dict__) + exec( + compile(ast.Module(body=[assignment, validator], type_ignores=[]), str(SGLANG_ROLLOUT_PATH), "exec"), + module.__dict__, + ) return module._validate_p3o_behavior_sampling_params @@ -63,11 +69,11 @@ def test_p3o_sampling_contract_leaves_evaluation_and_non_p3o_unchanged(): def test_p3o_dispatch_requires_an_explicit_custom_generation_contract(): tree = ast.parse(SGLANG_ROLLOUT_PATH.read_text(encoding="utf-8")) - dispatch = next(node for node in tree.body if isinstance(node, ast.AsyncFunctionDef) and node.name == "_dispatch_generate") + dispatch = next( + node for node in tree.body if isinstance(node, ast.AsyncFunctionDef) and node.name == "_dispatch_generate" + ) string_constants = { - node.value - for node in ast.walk(dispatch) - if isinstance(node, ast.Constant) and isinstance(node.value, str) + node.value for node in ast.walk(dispatch) if isinstance(node, ast.Constant) and isinstance(node.value, str) } assert "p3o_behavior_logprob_contract" in string_constants From 3cdf4a98611f40b50c716887b205a36b0a7196e4 Mon Sep 17 00:00:00 2001 From: Chream Date: Fri, 14 Aug 2026 11:10:21 +0800 Subject: [PATCH 06/30] test: preserve import cache in P3O tests --- tests/backends/megatron/_megatron_stub.py | 22 +++++++++++++++ .../backends/megatron/test_p3o_distributed.py | 6 ++--- .../backends/megatron/test_p3o_model_step.py | 27 ++++++++++++++++--- 3 files changed, 48 insertions(+), 7 deletions(-) diff --git a/tests/backends/megatron/_megatron_stub.py b/tests/backends/megatron/_megatron_stub.py index 350f04ee8..0699e5282 100644 --- a/tests/backends/megatron/_megatron_stub.py +++ b/tests/backends/megatron/_megatron_stub.py @@ -32,6 +32,28 @@ #: submodule below these is synthesized on demand, so the P3O import chain does #: not have to be enumerated here. STUBBED_ROOTS = ("megatron",) +_MISSING = object() + + +@contextmanager +def temporarily_stub_module(name: str, module: ModuleType) -> Iterator[None]: + """Expose one import stub without resetting unrelated import-cache entries. + + ``unittest.mock.patch.dict(sys.modules, ...)`` restores the complete module + cache when its context exits. Imports performed inside that context can + include native Torch or Ray modules, so clearing them makes a later import + reinitialize process-global extension state. Preserve and restore only the + requested entry instead. + """ + previous_module = sys.modules.get(name, _MISSING) + sys.modules[name] = module + try: + yield + finally: + if previous_module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous_module class _MagicModule(ModuleType): diff --git a/tests/backends/megatron/test_p3o_distributed.py b/tests/backends/megatron/test_p3o_distributed.py index b86f245b7..832635aa2 100644 --- a/tests/backends/megatron/test_p3o_distributed.py +++ b/tests/backends/megatron/test_p3o_distributed.py @@ -7,9 +7,7 @@ import math import os import socket -import sys from types import ModuleType -from unittest.mock import patch import torch import torch.distributed as dist @@ -21,14 +19,14 @@ compute_p3o_token_terms, finalize_p3o_step_context, ) -from tests.backends.megatron._megatron_stub import stubbed_megatron_modules +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules, temporarily_stub_module stream_dataloader = ModuleType("relax.utils.data.stream_dataloader") stream_dataloader.StreamingTQIterator = object with ( - patch.dict(sys.modules, {"relax.utils.data.stream_dataloader": stream_dataloader}), + temporarily_stub_module("relax.utils.data.stream_dataloader", stream_dataloader), stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), ): from relax.backends.megatron import loss as loss_module diff --git a/tests/backends/megatron/test_p3o_model_step.py b/tests/backends/megatron/test_p3o_model_step.py index 88dbd436c..054e7c561 100644 --- a/tests/backends/megatron/test_p3o_model_step.py +++ b/tests/backends/megatron/test_p3o_model_step.py @@ -9,12 +9,11 @@ from argparse import Namespace from pathlib import Path from types import ModuleType, SimpleNamespace -from unittest.mock import patch import pytest import torch -from tests.backends.megatron._megatron_stub import stubbed_megatron_modules +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules, temporarily_stub_module MODEL_PATH = Path(__file__).resolve().parents[3] / "relax" / "backends" / "megatron" / "model.py" @@ -23,7 +22,7 @@ stream_dataloader.StreamingTQIterator = object with ( - patch.dict(sys.modules, {"relax.utils.data.stream_dataloader": stream_dataloader}), + temporarily_stub_module("relax.utils.data.stream_dataloader", stream_dataloader), stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), ): from relax.backends.megatron import model as model_module @@ -32,6 +31,28 @@ _preserved_dynamic_cp_group = model_module._preserved_dynamic_cp_group +def test_p3o_module_stub_preserves_unrelated_import_cache(monkeypatch): + stub_name = "tests.backends.megatron._p3o_stream_dataloader_stub" + unrelated_name = "tests.backends.megatron._p3o_import_cache_probe" + stub_module = ModuleType(stub_name) + unrelated_module = ModuleType(unrelated_name) + monkeypatch.delitem(sys.modules, stub_name, raising=False) + monkeypatch.delitem(sys.modules, unrelated_name, raising=False) + + with temporarily_stub_module(stub_name, stub_module): + sys.modules[unrelated_name] = unrelated_module + + assert stub_name not in sys.modules + assert sys.modules[unrelated_name] is unrelated_module + + previous_module = ModuleType(stub_name) + monkeypatch.setitem(sys.modules, stub_name, previous_module) + with temporarily_stub_module(stub_name, stub_module): + assert sys.modules[stub_name] is stub_module + + assert sys.modules[stub_name] is previous_module + + def test_p3o_model_step_restores_dynamic_cp_group_after_error(): original_group = object() dynamic_group = object() From 0567376132771f13617e55fd8ef8ff0636209459 Mon Sep 17 00:00:00 2001 From: Chream Date: Fri, 14 Aug 2026 11:21:35 +0800 Subject: [PATCH 07/30] test: isolate P3O Megatron stubs --- tests/backends/megatron/_megatron_stub.py | 29 +++++++++++++++++ .../backends/megatron/test_p3o_distributed.py | 7 +++- .../backends/megatron/test_p3o_model_step.py | 32 ++++++++++++++++++- 3 files changed, 66 insertions(+), 2 deletions(-) diff --git a/tests/backends/megatron/_megatron_stub.py b/tests/backends/megatron/_megatron_stub.py index 0699e5282..899a718c5 100644 --- a/tests/backends/megatron/_megatron_stub.py +++ b/tests/backends/megatron/_megatron_stub.py @@ -56,6 +56,35 @@ def temporarily_stub_module(name: str, module: ModuleType) -> Iterator[None]: sys.modules[name] = previous_module +@contextmanager +def isolated_module_cache(prefix: str) -> Iterator[None]: + """Discard test-only imports below ``prefix`` without flushing all + modules.""" + prefix_with_dot = f"{prefix}." + modules_before = { + name: module for name, module in sys.modules.items() if name == prefix or name.startswith(prefix_with_dot) + } + try: + yield + finally: + new_module_names = sorted( + ( + name + for name in sys.modules + if name not in modules_before and (name == prefix or name.startswith(prefix_with_dot)) + ), + key=lambda module_name: module_name.count("."), + reverse=True, + ) + for name in new_module_names: + module = sys.modules.pop(name) + parent_name, _, attribute_name = name.rpartition(".") + parent_module = sys.modules.get(parent_name) + if parent_module is not None and getattr(parent_module, attribute_name, _MISSING) is module: + delattr(parent_module, attribute_name) + sys.modules.update(modules_before) + + class _MagicModule(ModuleType): """Module whose unknown attributes resolve to ``MagicMock``. diff --git a/tests/backends/megatron/test_p3o_distributed.py b/tests/backends/megatron/test_p3o_distributed.py index 832635aa2..0bad98585 100644 --- a/tests/backends/megatron/test_p3o_distributed.py +++ b/tests/backends/megatron/test_p3o_distributed.py @@ -19,13 +19,18 @@ compute_p3o_token_terms, finalize_p3o_step_context, ) -from tests.backends.megatron._megatron_stub import stubbed_megatron_modules, temporarily_stub_module +from tests.backends.megatron._megatron_stub import ( + isolated_module_cache, + stubbed_megatron_modules, + temporarily_stub_module, +) stream_dataloader = ModuleType("relax.utils.data.stream_dataloader") stream_dataloader.StreamingTQIterator = object with ( + isolated_module_cache("relax.backends.megatron"), temporarily_stub_module("relax.utils.data.stream_dataloader", stream_dataloader), stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), ): diff --git a/tests/backends/megatron/test_p3o_model_step.py b/tests/backends/megatron/test_p3o_model_step.py index 054e7c561..6723967ce 100644 --- a/tests/backends/megatron/test_p3o_model_step.py +++ b/tests/backends/megatron/test_p3o_model_step.py @@ -13,7 +13,11 @@ import pytest import torch -from tests.backends.megatron._megatron_stub import stubbed_megatron_modules, temporarily_stub_module +from tests.backends.megatron._megatron_stub import ( + isolated_module_cache, + stubbed_megatron_modules, + temporarily_stub_module, +) MODEL_PATH = Path(__file__).resolve().parents[3] / "relax" / "backends" / "megatron" / "model.py" @@ -22,6 +26,7 @@ stream_dataloader.StreamingTQIterator = object with ( + isolated_module_cache("relax.backends.megatron"), temporarily_stub_module("relax.utils.data.stream_dataloader", stream_dataloader), stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), ): @@ -53,6 +58,31 @@ def test_p3o_module_stub_preserves_unrelated_import_cache(monkeypatch): assert sys.modules[stub_name] is previous_module +def test_p3o_module_cache_scope_removes_only_test_imports(monkeypatch): + prefix = "tests.backends.megatron._p3o_module_cache" + child_name = f"{prefix}.child" + unrelated_name = "tests.backends.megatron._p3o_unrelated_import" + cache_module = ModuleType(prefix) + child_module = ModuleType(child_name) + unrelated_module = ModuleType(unrelated_name) + parent_module = sys.modules["tests.backends.megatron"] + monkeypatch.delitem(sys.modules, prefix, raising=False) + monkeypatch.delitem(sys.modules, child_name, raising=False) + monkeypatch.delattr(parent_module, "_p3o_module_cache", raising=False) + + with isolated_module_cache(prefix): + sys.modules[prefix] = cache_module + setattr(parent_module, "_p3o_module_cache", cache_module) + sys.modules[child_name] = child_module + setattr(cache_module, "child", child_module) + monkeypatch.setitem(sys.modules, unrelated_name, unrelated_module) + + assert prefix not in sys.modules + assert child_name not in sys.modules + assert not hasattr(parent_module, "_p3o_module_cache") + assert sys.modules[unrelated_name] is unrelated_module + + def test_p3o_model_step_restores_dynamic_cp_group_after_error(): original_group = object() dynamic_group = object() From 7e11a492ee281a1c34846bc51d8fb84eb2474685 Mon Sep 17 00:00:00 2001 From: Chream Date: Fri, 14 Aug 2026 11:34:43 +0800 Subject: [PATCH 08/30] test: wait for reward worker teardown --- tests/engine/rewards/test_custom_reward_worker.py | 4 ++-- tests/engine/rewards/test_reward_router.py | 13 ++++++------- tests/engine/rewards/test_reward_worker.py | 13 ++++++------- 3 files changed, 14 insertions(+), 16 deletions(-) diff --git a/tests/engine/rewards/test_custom_reward_worker.py b/tests/engine/rewards/test_custom_reward_worker.py index 3805b1f19..a81f19972 100644 --- a/tests/engine/rewards/test_custom_reward_worker.py +++ b/tests/engine/rewards/test_custom_reward_worker.py @@ -90,8 +90,8 @@ def _kill_workers(executor: RewardExecutor | None) -> None: return for worker in executor._workers: try: - ray.kill(worker) - except Exception: + ray.get(worker.__ray_terminate__.remote()) + except ray.exceptions.RayError: # A worker may already be dead after an expected failure. Teardown # must still clear the remaining actor handles. pass diff --git a/tests/engine/rewards/test_reward_router.py b/tests/engine/rewards/test_reward_router.py index 5fd2b5a65..c647372ff 100644 --- a/tests/engine/rewards/test_reward_router.py +++ b/tests/engine/rewards/test_reward_router.py @@ -68,20 +68,19 @@ def _make_sample(response: str, label: str, metadata: dict | None = None) -> "Sa def _kill_executor_workers(): - """Kill named Ray actors held by the current RewardExecutor singleton. + """Synchronously terminate named actors held by the executor singleton. - Without this, dropping ``RewardExecutor._instance`` only releases the - Python actor handles; on Python 3.10 the next test can reach - ``options(get_if_exists=True)`` before Ray finishes evicting the named - actors, get a handle to a dying actor, and hit ActorDiedError. + ``ray.kill`` is asynchronous, so the next test can reach + ``options(get_if_exists=True)`` before Ray evicts a same-named actor and + receive a handle that is about to die. """ inst = RewardExecutor._instance if inst is None: return for w in inst._workers: try: - ray.kill(w) - except Exception: + ray.get(w.__ray_terminate__.remote()) + except ray.exceptions.RayError: pass inst._workers = [] diff --git a/tests/engine/rewards/test_reward_worker.py b/tests/engine/rewards/test_reward_worker.py index 38040222e..9f36208ea 100644 --- a/tests/engine/rewards/test_reward_worker.py +++ b/tests/engine/rewards/test_reward_worker.py @@ -244,20 +244,19 @@ def test_unknown_rm_type_raises(self): def _kill_executor_workers(): - """Kill named Ray actors held by the current RewardExecutor singleton. + """Synchronously terminate named actors held by the executor singleton. - Without this, dropping ``RewardExecutor._instance`` only releases the - Python actor handles; on Python 3.10 the next test can reach - ``options(get_if_exists=True)`` before Ray finishes evicting the named - actors, get a handle to a dying actor, and hit ActorDiedError. + ``ray.kill`` is asynchronous, so the next test can reach + ``options(get_if_exists=True)`` before Ray evicts a same-named actor and + receive a handle that is about to die. """ inst = RewardExecutor._instance if inst is None: return for w in inst._workers: try: - ray.kill(w) - except Exception: + ray.get(w.__ray_terminate__.remote()) + except ray.exceptions.RayError: pass inst._workers = [] From 6c7a3d2c0cd9e714d96a15e52a7d542a422230b5 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Fri, 14 Aug 2026 11:56:37 +0800 Subject: [PATCH 09/30] fix(rewards): isolate worker actor pools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Prevent reuse of terminating reward workers - Create anonymous RewardWorker actors owned by each executor - Avoid reacquiring same-named actors during Ray teardown - Update teardown documentation for the private worker lifecycle --- relax/engine/rewards/__init__.py | 10 +--------- tests/engine/rewards/test_reward_router.py | 7 +------ tests/engine/rewards/test_reward_worker.py | 7 +------ 3 files changed, 3 insertions(+), 21 deletions(-) diff --git a/relax/engine/rewards/__init__.py b/relax/engine/rewards/__init__.py index 11391a7ff..901d7474d 100644 --- a/relax/engine/rewards/__init__.py +++ b/relax/engine/rewards/__init__.py @@ -225,15 +225,7 @@ async def _ensure_workers(self): async with self._worker_init_lock: if self._workers: return - self._workers = await asyncio.to_thread( - lambda: [ - RewardWorker.options( - name=f"reward_worker_{i}", - get_if_exists=True, - ).remote() - for i in range(self._num_workers) - ] - ) + self._workers = await asyncio.to_thread(lambda: [RewardWorker.remote() for _ in range(self._num_workers)]) logger.info( "RewardExecutor: created %d RewardWorker actors (max_concurrency=%d)", self._num_workers, diff --git a/tests/engine/rewards/test_reward_router.py b/tests/engine/rewards/test_reward_router.py index c647372ff..fc8564212 100644 --- a/tests/engine/rewards/test_reward_router.py +++ b/tests/engine/rewards/test_reward_router.py @@ -68,12 +68,7 @@ def _make_sample(response: str, label: str, metadata: dict | None = None) -> "Sa def _kill_executor_workers(): - """Synchronously terminate named actors held by the executor singleton. - - ``ray.kill`` is asynchronous, so the next test can reach - ``options(get_if_exists=True)`` before Ray evicts a same-named actor and - receive a handle that is about to die. - """ + """Synchronously terminate actors held by the executor singleton.""" inst = RewardExecutor._instance if inst is None: return diff --git a/tests/engine/rewards/test_reward_worker.py b/tests/engine/rewards/test_reward_worker.py index 9f36208ea..bdd58cf8c 100644 --- a/tests/engine/rewards/test_reward_worker.py +++ b/tests/engine/rewards/test_reward_worker.py @@ -244,12 +244,7 @@ def test_unknown_rm_type_raises(self): def _kill_executor_workers(): - """Synchronously terminate named actors held by the executor singleton. - - ``ray.kill`` is asynchronous, so the next test can reach - ``options(get_if_exists=True)`` before Ray evicts a same-named actor and - receive a handle that is about to die. - """ + """Synchronously terminate actors held by the executor singleton.""" inst = RewardExecutor._instance if inst is None: return From cd020f101e41779216c740d5249266fc897fe8ff Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Fri, 14 Aug 2026 16:05:30 +0800 Subject: [PATCH 10/30] fix(p3o): harden distributed validation --- relax/backends/megatron/model_provider.py | 7 +++ relax/backends/megatron/p3o_step.py | 21 ++++--- relax/utils/training/p3o_utils.py | 55 ++++++++++++++----- .../megatron/test_model_provider_vpp.py | 11 ++++ .../backends/megatron/test_p3o_distributed.py | 29 +++++++--- tests/backends/megatron/test_p3o_step.py | 21 +++++++ tests/utils/training/test_p3o_utils.py | 34 ++++++++++++ 7 files changed, 146 insertions(+), 32 deletions(-) diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 7e9805608..0893b3366 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -313,6 +313,13 @@ def wrapped_model_provider( provider.fp16 = False provider.bf16 = True provider.params_dtype = torch.bfloat16 + else: + # The HF checkpoint may advertise a reduced-precision dtype. When + # both Megatron precision flags are disabled, make the requested + # FP32 mode explicit instead of inheriting that checkpoint hint. + provider.fp16 = False + provider.bf16 = False + provider.params_dtype = torch.float32 provider.finalize() diff --git a/relax/backends/megatron/p3o_step.py b/relax/backends/megatron/p3o_step.py index fbf3e677a..54951904b 100644 --- a/relax/backends/megatron/p3o_step.py +++ b/relax/backends/megatron/p3o_step.py @@ -46,7 +46,7 @@ P3O_STEP_CONTEXT_ATTR = "_p3o_step_context" P3O_NONFINITE_RATIO_ERROR = ( - "P3O: non-finite importance ratio at a valid response token on at least one rank; " + "P3O: non-finite importance ratio or unrepresentable squared ratio at a valid response token on at least one rank; " "refusing to silently fall back to ESS=1. Check rollout log-probs and mask alignment." ) @@ -119,14 +119,17 @@ def synchronize_p3o_stats( group_src=torch.distributed.get_world_size(group=pp_group) - 1, ) - valid = vector[3] <= 0 - if valid.device.type == "cpu": - if not bool(valid): - raise ValueError(P3O_NONFINITE_RATIO_ERROR) - else: - # Keep the accelerator hot path asynchronous. Every rank observes the - # globally reduced invalid flag, so they all fail consistently. - torch._assert_async(valid, P3O_NONFINITE_RATIO_ERROR) + # A collective can overflow even when every rank contributed finite local + # moments, so validate the reduced S1/S2/N in addition to the synchronized + # per-rank invalid flag before finalization can fall back to ESS=1. + valid = (vector[3] <= 0) & torch.isfinite(vector[:3]).all() + # This is one host check per optimizer step (or micro-batch ESS scope), + # after every rank has received the same reduced verdict. A CUDA + # ``torch._assert_async`` would turn this expected input error into a + # device-side assert, poison NCCL, and abort Ray workers instead of raising + # a catchable all-rank exception. + if not bool(valid): + raise ValueError(P3O_NONFINITE_RATIO_ERROR) return P3OSufficientStats.from_vector(vector[:3]) diff --git a/relax/utils/training/p3o_utils.py b/relax/utils/training/p3o_utils.py index 261fe4cca..dbb052b73 100644 --- a/relax/utils/training/p3o_utils.py +++ b/relax/utils/training/p3o_utils.py @@ -34,7 +34,7 @@ # Shared by the checked and unchecked sufficient-statistics paths so the message # a user sees does not depend on which one detected the non-finite ratio. NONFINITE_RATIO_MESSAGE = ( - "P3O: non-finite importance ratio at a valid response token; refusing to " + "P3O: non-finite importance ratio or unrepresentable squared ratio at a valid response token; refusing to " "silently fall back to ESS=1. Check rollout log-probs and mask alignment." ) @@ -188,7 +188,8 @@ def compute_p3o_sufficient_stats( Local :class:`P3OSufficientStats` in float64. Raises: - ValueError: If a valid position produced a non-finite ratio. + ValueError: If a valid position produced a non-finite ratio or a + squared ratio that overflowed or underflowed in float64. """ stats, invalid_flag = compute_p3o_sufficient_stats_unchecked(log_probs, behavior_log_probs, valid_mask) # This convenience wrapper is used outside the micro-batch hot path, so an @@ -221,9 +222,10 @@ def compute_p3o_sufficient_stats_unchecked( Returns: ``(stats, invalid_flag)``. ``invalid_flag`` is ``1.0`` when any valid - position produced a non-finite ratio, else ``0.0``. When it is set, the - statistics are zeroed so a caller that defers the check cannot poison - ``S1/S2`` with ``inf``/``nan`` in the meantime. + position produced a non-finite ratio or a squared ratio that cannot be + represented in float64, else ``0.0``. When it is set, the statistics + are zeroed so a caller that defers the check cannot poison ``S1/S2`` + with ``inf``/``nan`` in the meantime. """ with torch.no_grad(): mask_bool = valid_mask.bool() @@ -231,19 +233,44 @@ def compute_p3o_sufficient_stats_unchecked( ratio = torch.exp(log_ratio.to(torch.float64)) ratio = torch.where(mask_bool, ratio, torch.zeros_like(ratio)) - - # Both checks stay on device. log_ratio is already zeroed outside the - # mask, so a global isfinite() over it is equivalent to masking first. - invalid_flag = (~(torch.isfinite(log_ratio).all() & torch.isfinite(ratio).all())).to(torch.float64) + ratio_sq = ratio.square() + + # All checks stay on device. A finite ratio can still overflow when + # squared (for example exp(500)), while a very small finite ratio can + # produce an exact-zero square (for example exp(-500)). Either would + # corrupt S2 and must fail through the synchronized invalid flag. Zero + # squares outside the valid-token mask remain harmless. + ratio_sq_representable = torch.where(mask_bool, ratio_sq > 0.0, torch.ones_like(mask_bool)).all() + token_invalid = ( + ~( + torch.isfinite(log_ratio).all() + & torch.isfinite(ratio).all() + & torch.isfinite(ratio_sq).all() + & ratio_sq_representable + ) + ).to(torch.float64) + + # Per-token finiteness is not sufficient: summing several large but + # finite squares can overflow locally before the collective. + local_sums = torch.stack( + ( + ratio.sum(), + ratio_sq.sum(), + mask_bool.sum().to(torch.float64), + ) + ) + invalid_flag = torch.maximum(token_invalid, (~torch.isfinite(local_sums).all()).to(torch.float64)) # Zero the contribution when invalid, so deferring the host-side check - # cannot let inf/nan reach the reduced moments. - keep = 1.0 - invalid_flag + # cannot let inf/nan reach the reduced moments. torch.where is required + # here: multiplying an infinite S2 by zero would still produce NaN. + keep = invalid_flag <= 0.0 + safe_local_sums = torch.where(keep, local_sums, torch.zeros_like(local_sums)) return ( P3OSufficientStats( - sum_ratio=ratio.sum() * keep, - sum_ratio_sq=ratio.pow(2).sum() * keep, - valid_token_count=mask_bool.sum().to(torch.float64) * keep, + sum_ratio=safe_local_sums[0], + sum_ratio_sq=safe_local_sums[1], + valid_token_count=safe_local_sums[2], ), invalid_flag, ) diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 82c20510b..9a273ebec 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -183,6 +183,17 @@ def test_bridge_provider_receives_virtual_pipeline_size(monkeypatch): assert provider.calls == [{"pre_process": True, "post_process": False, "vp_stage": 1}] +def test_bridge_provider_uses_fp32_when_precision_flags_are_disabled(monkeypatch): + module, provider = _load_model_provider(monkeypatch) + + module.get_model_provider_func(_bridge_args(fp16=False, bf16=False), role="actor") + + assert provider.fp16 is False + assert provider.bf16 is False + assert provider.params_dtype is torch.float32 + assert provider.finalized + + def test_bridge_provider_receives_vision_dp_when_cp(monkeypatch): module, provider = _load_model_provider(monkeypatch) diff --git a/tests/backends/megatron/test_p3o_distributed.py b/tests/backends/megatron/test_p3o_distributed.py index 0bad98585..c69a12c3a 100644 --- a/tests/backends/megatron/test_p3o_distributed.py +++ b/tests/backends/megatron/test_p3o_distributed.py @@ -16,6 +16,7 @@ from relax.utils.training.p3o_utils import ( P3OSufficientStats, compute_p3o_sufficient_stats, + compute_p3o_sufficient_stats_unchecked, compute_p3o_token_terms, finalize_p3o_step_context, ) @@ -52,19 +53,24 @@ def _init_gloo(rank: int, world_size: int, port: int) -> None: dist.init_process_group("gloo", rank=rank, world_size=world_size) -def _nonfinite_worker(rank: int, world_size: int, port: int) -> None: +def _extreme_ratio_worker(rank: int, world_size: int, port: int, log_ratio: float) -> None: _init_gloo(rank, world_size, port) try: p3o_step.mpu.is_pipeline_last_stage = lambda ignore_virtual=True: True p3o_step.mpu.get_data_parallel_group = lambda with_context_parallel=True: dist.group.WORLD p3o_step.mpu.get_pipeline_model_parallel_world_size = lambda: 1 - stats = ( - P3OSufficientStats.zeros() - if rank == 0 - else P3OSufficientStats.from_vector(torch.tensor([1.0, 1.0, 1.0], dtype=torch.float64)) - ) - invalid_count = torch.tensor(float(rank == 0), dtype=torch.float64) + if rank == 0: + stats, invalid_count = compute_p3o_sufficient_stats_unchecked( + torch.tensor([log_ratio], dtype=torch.float32), + torch.zeros(1, dtype=torch.float32), + torch.ones(1, dtype=torch.bool), + ) + else: + # Model a dummy micro-batch: it contributes neither moments nor an + # invalid flag, but must still observe the real rank's failure. + stats = P3OSufficientStats.zeros() + invalid_count = torch.zeros((), dtype=torch.float64) try: synchronize_p3o_stats( @@ -319,9 +325,14 @@ def _real_dummy_dp_worker(rank: int, world_size: int, port: int) -> None: dist.destroy_process_group() -def test_p3o_distributed_nonfinite_fails_synchronously(): +def test_p3o_distributed_positive_extreme_ratio_fails_synchronously_with_dummy_rank(): + world_size = 2 + mp.spawn(_extreme_ratio_worker, args=(world_size, _free_port(), 500.0), nprocs=world_size, join=True) + + +def test_p3o_distributed_negative_extreme_ratio_fails_synchronously_with_dummy_rank(): world_size = 2 - mp.spawn(_nonfinite_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) + mp.spawn(_extreme_ratio_worker, args=(world_size, _free_port(), -500.0), nprocs=world_size, join=True) def test_p3o_distributed_pipeline_broadcasts_last_stage_stats(): diff --git a/tests/backends/megatron/test_p3o_step.py b/tests/backends/megatron/test_p3o_step.py index b131a189a..7f75340ba 100644 --- a/tests/backends/megatron/test_p3o_step.py +++ b/tests/backends/megatron/test_p3o_step.py @@ -104,6 +104,27 @@ def all_reduce(vector, *, op, group): ) +def test_p3o_step_raises_when_collective_overflows_finite_local_stats(monkeypatch): + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + dp_cp_group = object() + + def all_reduce(vector, *, op, group): + assert group is dp_cp_group + vector[1] = float("inf") + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + with pytest.raises(ValueError, match="unrepresentable squared ratio"): + synchronize_p3o_stats( + _stats((1.0, torch.finfo(torch.float64).max * 0.75, 1.0)), + torch.zeros((), dtype=torch.float64), + dp_cp_group=dp_cp_group, + pp_group=None, + is_pipeline_last_stage=True, + ) + + def test_compute_p3o_step_context_plain_text_forward_kwargs(monkeypatch): """ESS pre-pass forward_step must use tokens+packed_seq_params for plain text.""" diff --git a/tests/utils/training/test_p3o_utils.py b/tests/utils/training/test_p3o_utils.py index ebbd87820..e89419f72 100644 --- a/tests/utils/training/test_p3o_utils.py +++ b/tests/utils/training/test_p3o_utils.py @@ -18,6 +18,7 @@ compute_p3o_behavior_kl_proxy, compute_p3o_exact_kl, compute_p3o_sufficient_stats, + compute_p3o_sufficient_stats_unchecked, compute_p3o_token_terms, finalize_p3o_step_context, ) @@ -265,6 +266,39 @@ def test_p3o_utils_non_finite_valid_token_raises(log_prob, behavior_log_prob): compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) +@pytest.mark.parametrize("log_ratio", [500.0, -500.0]) +def test_p3o_utils_extreme_finite_log_ratio_flags_moment_overflow_or_underflow(log_ratio): + behavior_log_probs = torch.zeros(1, 2, dtype=torch.float32) + log_probs = torch.tensor([[log_ratio, 0.0]], dtype=torch.float32) + valid_mask = torch.tensor([[True, False]]) + + stats, invalid_flag = compute_p3o_sufficient_stats_unchecked( + log_probs, + behavior_log_probs, + valid_mask, + ) + + assert float(invalid_flag) == 1.0 + torch.testing.assert_close(stats.as_vector(), torch.zeros(3, dtype=torch.float64), rtol=0.0, atol=0.0) + with pytest.raises(ValueError, match="non-finite importance ratio"): + compute_p3o_sufficient_stats(log_probs, behavior_log_probs, valid_mask) + + +def test_p3o_utils_finite_per_token_squares_reject_overflowing_local_sum(): + behavior_log_probs = torch.zeros(2, dtype=torch.float32) + log_probs = torch.full((2,), 354.8, dtype=torch.float32) + valid_mask = torch.ones(2, dtype=torch.bool) + + stats, invalid_flag = compute_p3o_sufficient_stats_unchecked( + log_probs, + behavior_log_probs, + valid_mask, + ) + + assert float(invalid_flag) == 1.0 + torch.testing.assert_close(stats.as_vector(), torch.zeros(3, dtype=torch.float64), rtol=0.0, atol=0.0) + + def test_p3o_utils_all_masked_poison_produces_fp64_zero_stats(): log_probs = torch.tensor([[float("nan"), float("inf")]], dtype=torch.float32) behavior_log_probs = torch.tensor([[float("-inf"), float("nan")]], dtype=torch.float32) From be888d2e1ebfa8d27b9ecbf51a482a9222cfac5b Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:04:08 +0800 Subject: [PATCH 11/30] chore(p3o): add offline oracle audit script --- .../local_p3o_task40/audit_oracle_vectors.py | 538 ++++++++++++++++++ 1 file changed, 538 insertions(+) create mode 100644 scripts/local_p3o_task40/audit_oracle_vectors.py diff --git a/scripts/local_p3o_task40/audit_oracle_vectors.py b/scripts/local_p3o_task40/audit_oracle_vectors.py new file mode 100644 index 000000000..f65f95197 --- /dev/null +++ b/scripts/local_p3o_task40/audit_oracle_vectors.py @@ -0,0 +1,538 @@ +#!/usr/bin/env python3 + +"""Audit Task40 Step-0 oracle inputs across DP4CP1 and DP2CP2.""" + +import argparse +import hashlib +import json +import math +from collections import defaultdict +from pathlib import Path +from typing import Any + +import torch + + +QUANTILES = (0.0, 0.5, 0.9, 0.95, 0.99, 1.0) +EXPECTED_TOPOLOGY_CONFIG_FIELDS = {"context_parallel_size"} +PRECISION_CONFIG_FIELDS = ("fp16", "bf16", "params_dtype", "autocast_dtype", "attention_softmax_in_fp32") +PARALLEL_CONFIG_FIELDS = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "context_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", +) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _token_key(tokens: torch.Tensor) -> str: + return hashlib.sha256(tokens.to(torch.int64).numpy().tobytes()).hexdigest() + + +def _as_float(value: Any) -> float: + if isinstance(value, torch.Tensor): + return float(value) + return float(value) + + +def _distribution(values: list[float]) -> dict[str, Any]: + if not values: + return {"count": 0, "max_abs": None, "quantiles": {}} + tensor = torch.tensor(values, dtype=torch.float64).abs() + return { + "count": tensor.numel(), + "max_abs": float(tensor.max()), + "mean_abs": float(tensor.mean()), + "quantiles": { + f"p{int(quantile * 100):02d}": float(torch.quantile(tensor, quantile)) for quantile in QUANTILES + }, + } + + +def _response_indices( + *, total_length: int, response_length: int, cp_size: int, cp_rank: int, max_seq_len: int +) -> list[tuple[int, int]]: + """Return ``(response_index, global_chunk_index)`` in local output order.""" + if cp_size == 1: + return [(index, 0) for index in range(response_length)] + + prompt_length = total_length - response_length + chunk_size = math.ceil(max_seq_len / (2 * cp_size)) + chunks = (cp_rank, 2 * cp_size - cp_rank - 1) + result = [] + for chunk_index in chunks: + # Logits at sequence index i predict token i+1. Match + # get_logits_and_tokens_offset_with_cp rather than slicing token IDs + # directly, including the one-token shift at every CP chunk boundary. + start = max(chunk_index * chunk_size, prompt_length - 1) + 1 + stop = min((chunk_index + 1) * chunk_size, total_length - 1) + 1 + for token_index in range(start, max(start, stop)): + if prompt_length <= token_index < total_length: + result.append((token_index - prompt_length, chunk_index)) + return result + + +def _runtime_by_rank(run_dir: Path) -> dict[int, dict[str, Any]]: + records = [json.loads(path.read_text()) for path in sorted((run_dir / "oracle").glob("runtime_rank*.json"))] + if not records: + raise ValueError(f"no runtime_rank*.json files under {run_dir / 'oracle'}") + return {int(record["rank"]): record for record in records} + + +def _new_sample( + tokens: torch.Tensor, + positions: torch.Tensor, + loss_mask: torch.Tensor, + total_length: int, + response_length: int, + max_seq_len: int, + max_seq_len_source: str, +) -> dict[str, Any]: + return { + "tokens": tokens.to(torch.int64), + "positions": positions.to(torch.int64), + "loss_mask": loss_mask.to(torch.int64), + "total_length": total_length, + "response_length": response_length, + "max_seq_len": max_seq_len, + "max_seq_len_source": max_seq_len_source, + "current": {}, + "rollout": {}, + "valid": {}, + "chunk": {}, + "cp_ranks": set(), + } + + +def _metadata_mismatch(sample: dict[str, Any], candidate: dict[str, Any]) -> str | None: + tensor_fields = ("tokens", "positions", "loss_mask") + for field in tensor_fields: + if not torch.equal(sample[field], candidate[field]): + return field + for field in ("total_length", "response_length", "max_seq_len"): + if sample[field] != candidate[field]: + return field + return None + + +def _load_run(run_dir: Path) -> dict[str, Any]: + runtimes = _runtime_by_rank(run_dir) + paths = sorted((run_dir / "oracle").glob("vectors_rank*_micro*.pt")) + if not paths: + raise ValueError(f"no vectors_rank*_micro*.pt files under {run_dir / 'oracle'}") + + samples: dict[str, dict[str, Any]] = {} + stats_by_rank: dict[int, dict[str, float]] = defaultdict(lambda: {"s1": 0.0, "s2": 0.0, "n": 0.0}) + duplicate_metadata_checks = 0 + vector_token_count = 0 + for path in paths: + artifact = torch.load(path, map_location="cpu", weights_only=False) + rank = int(artifact["rank"]) + runtime = runtimes[rank] + cp_rank = int(runtime["cp_rank"]) + cp_size = int(runtime["cp_world_size"]) + current = artifact["current_log_probs"].to(torch.float64).flatten() + rollout = artifact["rollout_log_probs"].to(torch.float64).flatten() + valid = artifact["valid_mask"].bool().flatten() + offset = 0 + raw_max_seq_lens = artifact.get("max_seq_lens") + for sample_index, (tokens, positions, loss_mask, total_length, response_length) in enumerate( + zip( + artifact["token_ids"], + artifact["position_ids"], + artifact["loss_masks"], + artifact["total_lengths"], + artifact["response_lengths"], + strict=True, + ) + ): + total_length = int(total_length) + response_length = int(response_length) + if raw_max_seq_lens is None: + max_seq_len = total_length + max_seq_len_source = "derived_from_total_length_for_thd_single_sample" + else: + max_seq_len = int(raw_max_seq_lens[sample_index]) + max_seq_len_source = "artifact" + key = _token_key(tokens) + candidate = _new_sample( + tokens, positions, loss_mask, total_length, response_length, max_seq_len, max_seq_len_source + ) + if key not in samples: + samples[key] = candidate + else: + mismatch = _metadata_mismatch(samples[key], candidate) + if mismatch is not None: + raise ValueError(f"{path}: duplicate sample {key} disagrees in {mismatch}") + duplicate_metadata_checks += 1 + sample = samples[key] + local_indices = _response_indices( + total_length=total_length, + response_length=response_length, + cp_size=cp_size, + cp_rank=cp_rank, + max_seq_len=max_seq_len, + ) + stop = offset + len(local_indices) + if stop > current.numel(): + raise ValueError(f"{path}: local vector ends inside sample {key}") + for local_index, (response_index, chunk_index) in enumerate(local_indices, start=offset): + if response_index in sample["current"]: + raise ValueError(f"{path}: duplicate response index {response_index} for sample {key}") + sample["current"][response_index] = float(current[local_index]) + sample["rollout"][response_index] = float(rollout[local_index]) + sample["valid"][response_index] = bool(valid[local_index]) + sample["chunk"][response_index] = chunk_index + sample["cp_ranks"].add(cp_rank) + offset = stop + vector_token_count += len(local_indices) + if not (offset == current.numel() == rollout.numel() == valid.numel()): + raise ValueError( + f"{path}: reconstructed={offset}, current={current.numel()}, rollout={rollout.numel()}, " + f"valid={valid.numel()}" + ) + stats_by_rank[rank]["s1"] += _as_float(artifact["local_s1"]) + stats_by_rank[rank]["s2"] += _as_float(artifact["local_s2"]) + stats_by_rank[rank]["n"] += _as_float(artifact["local_n"]) + + cp_size = max(int(record["cp_world_size"]) for record in runtimes.values()) + reconstruction_failures = [] + position_failures = [] + for key, sample in samples.items(): + expected_response_indices = set(range(sample["response_length"])) + if set(sample["current"]) != expected_response_indices: + reconstruction_failures.append( + { + "sample_key": key, + "missing": sorted(expected_response_indices - set(sample["current"])), + "unexpected": sorted(set(sample["current"]) - expected_response_indices), + } + ) + expected_positions = torch.arange(sample["total_length"], dtype=torch.int64) + if not torch.equal(sample["positions"], expected_positions): + position_failures.append(key) + + if cp_size > 1: + chunk_size = math.ceil(sample["max_seq_len"] / (2 * cp_size)) + reconstructed = torch.cat( + [ + sample["tokens"][chunk * chunk_size : min((chunk + 1) * chunk_size, sample["total_length"])] + for chunk in range(2 * cp_size) + ] + ) + if not torch.equal(reconstructed, sample["tokens"]): + reconstruction_failures.append({"sample_key": key, "reason": "ordered CP chunks != CP1 sequence"}) + + stats_sum = {field: sum(rank_stats[field] for rank_stats in stats_by_rank.values()) for field in ("s1", "s2", "n")} + runtime_contract = { + "bf16_all_ranks": all(record.get("bf16") is True for record in runtimes.values()), + "fp16_disabled_all_ranks": all(record.get("fp16") is False for record in runtimes.values()), + "requested_precisions": sorted({record.get("requested_precision") for record in runtimes.values()}), + "cp_world_sizes": sorted({int(record["cp_world_size"]) for record in runtimes.values()}), + "dp_world_sizes": sorted({int(record["dp_world_size"]) for record in runtimes.values()}), + "micro_batch_sizes": sorted({int(record["micro_batch_size"]) for record in runtimes.values()}), + "global_batch_sizes": sorted({int(record["global_batch_size"]) for record in runtimes.values()}), + "qkv_formats": sorted({record.get("qkv_format", "thd") for record in runtimes.values()}), + } + return { + "run_dir": str(run_dir.resolve()), + "runtime": list(runtimes.values()), + "vector_file_count": len(paths), + "sample_count": len(samples), + "vector_token_count": vector_token_count, + "duplicate_metadata_checks": duplicate_metadata_checks, + "samples": samples, + "position_global_monotonic": not position_failures, + "position_failure_sample_keys": position_failures, + "cp_reconstruction_pass": not reconstruction_failures, + "cp_reconstruction_failures": reconstruction_failures, + "runtime_contract": runtime_contract, + "stats_by_rank": {str(rank): values for rank, values in sorted(stats_by_rank.items())}, + "local_stats_sum": stats_sum, + } + + +def _first_tensor_difference(left: torch.Tensor, right: torch.Tensor) -> dict[str, Any] | None: + if left.shape != right.shape: + return {"index": 0, "left_shape": list(left.shape), "right_shape": list(right.shape)} + difference = (left != right).flatten().nonzero() + if difference.numel() == 0: + return None + index = int(difference[0]) + return {"index": index, "left": left.flatten()[index].item(), "right": right.flatten()[index].item()} + + +def _compare_inputs(left: dict[str, Any], right: dict[str, Any]) -> tuple[dict[str, Any], list[dict[str, Any]]]: + left_keys = set(left["samples"]) + right_keys = set(right["samples"]) + mismatches: list[dict[str, Any]] = [] + if left_keys != right_keys: + mismatches.append( + { + "field": "token_ids", + "left_only_sample_keys": sorted(left_keys - right_keys), + "right_only_sample_keys": sorted(right_keys - left_keys), + "first_difference_token": 0, + } + ) + + checks = {field: True for field in ("token_ids", "position_ids", "max_seq_lens", "loss_masks", "valid_masks")} + rollout_errors = [] + current_errors = [] + current_by_chunk: dict[int, list[float]] = defaultdict(list) + boundary_errors: dict[str, list[float]] = defaultdict(list) + valid_token_count = 0 + for key in sorted(left_keys & right_keys): + left_sample = left["samples"][key] + right_sample = right["samples"][key] + tensor_fields = {"token_ids": "tokens", "position_ids": "positions", "loss_masks": "loss_mask"} + for public_field, internal_field in tensor_fields.items(): + difference = _first_tensor_difference(left_sample[internal_field], right_sample[internal_field]) + if difference is not None: + checks[public_field] = False + mismatches.append( + { + "sample_key": key, + "field": public_field, + "first_difference_token": difference["index"], + **difference, + } + ) + left_valid = torch.tensor( + [left_sample["valid"][index] for index in sorted(left_sample["valid"])], dtype=torch.bool + ) + right_valid = torch.tensor( + [right_sample["valid"][index] for index in sorted(right_sample["valid"])], dtype=torch.bool + ) + valid_difference = _first_tensor_difference(left_valid, right_valid) + if valid_difference is not None: + checks["valid_masks"] = False + mismatches.append( + { + "sample_key": key, + "field": "valid_masks", + "first_difference_token": valid_difference["index"], + **valid_difference, + } + ) + if left_sample["max_seq_len"] != right_sample["max_seq_len"]: + checks["max_seq_lens"] = False + mismatches.append( + { + "sample_key": key, + "field": "max_seq_lens", + "first_difference_token": 0, + "left": left_sample["max_seq_len"], + "right": right_sample["max_seq_len"], + } + ) + + prompt_length = left_sample["total_length"] - left_sample["response_length"] + chunk_size = math.ceil(right_sample["max_seq_len"] / 4) + global_boundaries = [chunk_size, 2 * chunk_size, 3 * chunk_size] + for response_index in sorted(set(left_sample["current"]) & set(right_sample["current"])): + if not (left_sample["valid"][response_index] and right_sample["valid"][response_index]): + continue + valid_token_count += 1 + rollout_error = left_sample["rollout"][response_index] - right_sample["rollout"][response_index] + current_error = left_sample["current"][response_index] - right_sample["current"][response_index] + rollout_errors.append(rollout_error) + current_errors.append(current_error) + chunk_index = right_sample["chunk"][response_index] + current_by_chunk[chunk_index].append(current_error) + global_token_index = prompt_length + response_index + for boundary in global_boundaries: + distance = abs(global_token_index - boundary) + if distance <= 2: + boundary_errors[f"distance_{distance}"].append(current_error) + + checks["token_ids"] &= left_keys == right_keys + return ( + { + "field_equality": checks, + "valid_token_count": valid_token_count, + "rollout_log_probs": _distribution(rollout_errors), + "current_log_probs": { + "overall": _distribution(current_errors), + "by_cp_chunk": { + f"chunk_{chunk}": _distribution(values) for chunk, values in sorted(current_by_chunk.items()) + }, + "near_chunk_boundaries": { + name: _distribution(values) for name, values in sorted(boundary_errors.items()) + }, + }, + }, + mismatches, + ) + + +def _load_and_compare_configs(left_path: Path, right_path: Path) -> dict[str, Any]: + left = json.loads(left_path.read_text()) + right = json.loads(right_path.read_text()) + all_fields = sorted(set(left) | set(right)) + differences = { + field: {"dp4cp1": left.get(field), "dp2cp2": right.get(field)} + for field in all_fields + if left.get(field) != right.get(field) + } + unexpected = {field: value for field, value in differences.items() if field not in EXPECTED_TOPOLOGY_CONFIG_FIELDS} + return { + "dp4cp1_path": str(left_path.resolve()), + "dp2cp2_path": str(right_path.resolve()), + "dp4cp1_sha256": _sha256(left_path), + "dp2cp2_sha256": _sha256(right_path), + "field_count": len(all_fields), + "precision_fields": { + field: {"dp4cp1": left.get(field), "dp2cp2": right.get(field)} for field in PRECISION_CONFIG_FIELDS + }, + "parallel_fields": { + field: {"dp4cp1": left.get(field), "dp2cp2": right.get(field)} for field in PARALLEL_CONFIG_FIELDS + }, + "differences": differences, + "expected_topology_difference_fields": sorted(EXPECTED_TOPOLOGY_CONFIG_FIELDS), + "unexpected_differences": unexpected, + "invariant_fields_identical": not unexpected, + } + + +def _public_run(run: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in run.items() if key != "samples"} + + +def _render_markdown(report: dict[str, Any]) -> str: + comparison = report["comparison"] + config = report["provider_transformer_config"] + lines = [ + "# Task40 P3O offline input audit", + "", + f"**Verdict: `{report['verdict']}`**", + "", + "## Input equality", + "", + ] + for field, equal in comparison["field_equality"].items(): + lines.append(f"- `{field}`: {'equal' if equal else 'DIFFERENT'}") + lines.extend( + [ + f"- rollout log-prob max-abs: `{comparison['rollout_log_probs']['max_abs']}`", + f"- valid response tokens compared: `{comparison['valid_token_count']}`", + f"- DP2CP2 ordered chunk reconstruction: `{report['runs']['dp2cp2']['cp_reconstruction_pass']}`", + f"- global monotonic positions: DP4CP1=`{report['runs']['dp4cp1']['position_global_monotonic']}`, " + f"DP2CP2=`{report['runs']['dp2cp2']['position_global_monotonic']}`", + f"- BF16 runtime contract: DP4CP1=`{report['runs']['dp4cp1']['runtime_contract']['bf16_all_ranks']}`, " + f"DP2CP2=`{report['runs']['dp2cp2']['runtime_contract']['bf16_all_ranks']}`", + "", + "## Current log-prob error distribution", + "", + "```json", + json.dumps(comparison["current_log_probs"], indent=2, sort_keys=True), + "```", + "", + "## Local sufficient-stat sums", + "", + "```json", + json.dumps( + {name: run["local_stats_sum"] for name, run in report["runs"].items()}, indent=2, sort_keys=True + ), + "```", + "", + "## Provider transformer config", + "", + f"Compared `{config['field_count']}` top-level fields. Invariant fields identical: " + f"`{config['invariant_fields_identical']}`. The only allowed topology delta is " + "`context_parallel_size`.", + "", + "The BF16 oracle runs did not contain `provider_transformer_config.json`; these config files are from the " + "same P0 batch's final successful FP32 diagnostic runs. BF16 precision itself is confirmed by every " + "`runtime_rank*.json` record, but the provider-config comparison is explicitly a same-batch proxy.", + "", + "```json", + json.dumps( + { + "precision_fields": config["precision_fields"], + "parallel_fields": config["parallel_fields"], + "differences": config["differences"], + "unexpected_differences": config["unexpected_differences"], + }, + indent=2, + sort_keys=True, + ), + "```", + ] + ) + if report["mismatches"]: + lines.extend(["", "## First mismatch", "", "```json", json.dumps(report["mismatches"][0], indent=2), "```"]) + lines.extend( + [ + "", + "## Evidence limitation", + "", + "The THD vector hook stored full un-concatenated tokens and synthesized global position IDs, not the " + "post-slice packed token tensor or `cu_seqlens`. This audit validates the recorded inputs, inferred " + "single-sample `max_seq_lens`, CP ownership/coverage, and response-vector alignment; Batch 7 remains " + "responsible for direct `cu_seqlens` and weight-identity capture.", + "", + ] + ) + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dp4cp1-run", type=Path, required=True) + parser.add_argument("--dp2cp2-run", type=Path, required=True) + parser.add_argument("--dp4cp1-config", type=Path, required=True) + parser.add_argument("--dp2cp2-config", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + + runs = {"dp4cp1": _load_run(args.dp4cp1_run), "dp2cp2": _load_run(args.dp2cp2_run)} + comparison, mismatches = _compare_inputs(runs["dp4cp1"], runs["dp2cp2"]) + config = _load_and_compare_configs(args.dp4cp1_config, args.dp2cp2_config) + runtime_contract_pass = ( + runs["dp4cp1"]["runtime_contract"]["bf16_all_ranks"] + and runs["dp2cp2"]["runtime_contract"]["bf16_all_ranks"] + and runs["dp4cp1"]["runtime_contract"]["cp_world_sizes"] == [1] + and runs["dp4cp1"]["runtime_contract"]["dp_world_sizes"] == [4] + and runs["dp2cp2"]["runtime_contract"]["cp_world_sizes"] == [2] + and runs["dp2cp2"]["runtime_contract"]["dp_world_sizes"] == [2] + ) + input_equal = ( + all(comparison["field_equality"].values()) + and comparison["rollout_log_probs"]["max_abs"] == 0.0 + and all(run["position_global_monotonic"] and run["cp_reconstruction_pass"] for run in runs.values()) + and runtime_contract_pass + and config["invariant_fields_identical"] + ) + verdict = "INPUT_IDENTICAL" if input_equal else "INPUT_MISMATCH_FOUND" + report = { + "verdict": verdict, + "mismatches": mismatches, + "runs": {name: _public_run(run) for name, run in runs.items()}, + "comparison": comparison, + "provider_transformer_config": config, + "runtime_contract_pass": runtime_contract_pass, + "evidence_notes": { + "max_seq_lens": "derived from total_length because THD single-sample artifacts omit max_seq_lens", + "bf16_provider_config": "not captured; same-batch final successful FP32 configs used as proxy", + "actual_parameter_hash": "deferred to Batch 7 by plan", + "cu_seqlens": "not captured by this oracle hook; direct capture deferred to Batch 7", + }, + } + args.output_dir.mkdir(parents=True, exist_ok=True) + (args.output_dir / "offline_input_audit.json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + (args.output_dir / "offline_input_audit.md").write_text(_render_markdown(report)) + print(json.dumps({"verdict": verdict, "mismatch_count": len(mismatches)}, sort_keys=True)) + + +if __name__ == "__main__": + main() From a8b4eea0bc6cca62a13397b1e8c65d574a647bb4 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:26:03 +0800 Subject: [PATCH 12/30] fix(p3o): harden S2 overflow guards --- .../backends/megatron/test_p3o_distributed.py | 33 +++++++++++++++++++ tests/backends/megatron/test_p3o_step.py | 8 +++-- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/tests/backends/megatron/test_p3o_distributed.py b/tests/backends/megatron/test_p3o_distributed.py index c69a12c3a..15739dae7 100644 --- a/tests/backends/megatron/test_p3o_distributed.py +++ b/tests/backends/megatron/test_p3o_distributed.py @@ -92,6 +92,34 @@ def _extreme_ratio_worker(rank: int, world_size: int, port: int, log_ratio: floa dist.destroy_process_group() +def _reduced_s2_overflow_worker(rank: int, world_size: int, port: int) -> None: + _init_gloo(rank, world_size, port) + try: + local_stats = P3OSufficientStats.from_vector( + torch.tensor([1.0, torch.finfo(torch.float64).max * 0.75, 1.0], dtype=torch.float64) + ) + assert torch.isfinite(local_stats.as_vector()).all() + + try: + synchronize_p3o_stats( + local_stats, + torch.zeros((), dtype=torch.float64), + dp_cp_group=dist.group.WORLD, + pp_group=None, + is_pipeline_last_stage=True, + ) + except ValueError as error: + assert "unrepresentable squared ratio" in str(error) + else: + raise AssertionError("every rank must fail before the ESS=1 fallback") + + healthy = torch.ones((), dtype=torch.float64) + dist.all_reduce(healthy) + assert healthy.item() == world_size + finally: + dist.destroy_process_group() + + def _pipeline_worker(rank: int, world_size: int, port: int) -> None: _init_gloo(rank, world_size, port) try: @@ -335,6 +363,11 @@ def test_p3o_distributed_negative_extreme_ratio_fails_synchronously_with_dummy_r mp.spawn(_extreme_ratio_worker, args=(world_size, _free_port(), -500.0), nprocs=world_size, join=True) +def test_p3o_distributed_collective_s2_overflow_fails_synchronously(): + world_size = 2 + mp.spawn(_reduced_s2_overflow_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) + + def test_p3o_distributed_pipeline_broadcasts_last_stage_stats(): world_size = 2 mp.spawn(_pipeline_worker, args=(world_size, _free_port()), nprocs=world_size, join=True) diff --git a/tests/backends/megatron/test_p3o_step.py b/tests/backends/megatron/test_p3o_step.py index 7f75340ba..dab2b6b02 100644 --- a/tests/backends/megatron/test_p3o_step.py +++ b/tests/backends/megatron/test_p3o_step.py @@ -104,20 +104,22 @@ def all_reduce(vector, *, op, group): ) -def test_p3o_step_raises_when_collective_overflows_finite_local_stats(monkeypatch): +@pytest.mark.parametrize("moment_index", range(3), ids=("s1", "s2", "n")) +@pytest.mark.parametrize("nonfinite", [float("inf"), float("nan")], ids=("inf", "nan")) +def test_p3o_step_raises_when_collective_produces_nonfinite_moment(monkeypatch, moment_index, nonfinite): monkeypatch.setattr(torch.distributed, "is_available", lambda: True) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) dp_cp_group = object() def all_reduce(vector, *, op, group): assert group is dp_cp_group - vector[1] = float("inf") + vector[moment_index] = nonfinite monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) with pytest.raises(ValueError, match="unrepresentable squared ratio"): synchronize_p3o_stats( - _stats((1.0, torch.finfo(torch.float64).max * 0.75, 1.0)), + _stats((1.0, 1.0, 1.0)), torch.zeros((), dtype=torch.float64), dp_cp_group=dp_cp_group, pp_group=None, From 7dee79c9167a9a9068bfdbd44cd23570bcc12d54 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:53:12 +0800 Subject: [PATCH 13/30] feat(p3o): add CP forward diagnostics --- scripts/local_p3o_task40/README.md | 201 ++++++ .../local_p3o_task40/compare_forward_dumps.py | 485 +++++++++++++ scripts/local_p3o_task40/p3o_forward_dump.py | 666 ++++++++++++++++++ scripts/local_p3o_task40/run_forward_dump.sh | 123 ++++ 4 files changed, 1475 insertions(+) create mode 100644 scripts/local_p3o_task40/README.md create mode 100755 scripts/local_p3o_task40/compare_forward_dumps.py create mode 100755 scripts/local_p3o_task40/p3o_forward_dump.py create mode 100755 scripts/local_p3o_task40/run_forward_dump.sh diff --git a/scripts/local_p3o_task40/README.md b/scripts/local_p3o_task40/README.md new file mode 100644 index 000000000..67bea5745 --- /dev/null +++ b/scripts/local_p3o_task40/README.md @@ -0,0 +1,201 @@ +# Task40 P3O CP forward diagnostics + +This directory contains the Batch 3 harness for locating the first BF16 THD +forward divergence between CP1 and CP2. It is a diagnostic package only: it +does not change Relax production code and Batch 3 does not run the cluster +commands. + +## Provenance and deliberate differences + +The harness is an extension of two tools that were already exercised before +this batch; it is not a new independent implementation. + +1. The 2026-08-14 Step-0 oracle custom-init hook is + `scripts/local_p3o_task40/p0_oracle_hook.py` (SHA-256 + `d53fdded5f601e0611231520166850dd49f05aad754e36f8adcef0333d31bd8c`). + `p3o_forward_dump.py` retains its custom-init entry point, rank/DP/CP/TP/PP + runtime metadata, `dump_details`-relative output, detached CPU tensor + serialization, and SHA-256 key over the full int64 token sequence. The + oracle's P3O S1/S2/N wrapping and FP32 native-attention replacement are not + copied: this batch observes forward activations only and is BF16-only. +2. The module discovery and forward-hook technique comes from + `relax/backends/megatron/model_provider.py::_install_cp_probe` (the + `CP-PROBE` block). The new hook calls that existing probe first, then extends + its `named_modules()`/`with_kwargs=True` approach from one attention-shape + sample to every decoder layer and every captured log-prob micro-batch. It + writes tensors rather than log lines and adds global token ownership for the + CP zig-zag layout. Production `model_provider.py` is not modified. + +The two frozen 8/14 BF16 command sources used for parameter comparison are: + +- DP4CP1 command SHA-256: + `7e7fe652a059b9a88b0122a5fde7fa0c13a2d073b9cf667a7fc2eeea86895bae` +- DP2CP2 command SHA-256: + `fb5882f8b92d2b241c7910f842b32851a8379b2592feb1693b8cd74683c237e3` + +`run_forward_dump.sh` keeps their model, fixture, seed, rollout shape, global +batch, MBS1, optimizer, P3O step-scope settings, FlashAttention, no-recompute, +and deterministic rollout arguments. Its only training-argv additions are the +three diagnostic hook paths plus `--dump-details`; topology changes are limited +to `--resource` and `--context-parallel-size`, as shown below. The non-semantic +TensorBoard experiment name and output paths use the new topology/run names. +DP2CP1 uses the same 8/14 DP2 resource template (`actor=[1,2]`, +`rollout=[1,4]`) with the frozen Step-0 fixture. THD and BF16 are the same +resolved defaults used by the 8/14 BF16 commands and are enforced again by +`configure()`. + +| Template | Actor resource | Rollout resource | DP | CP | Precision | QKV | MBS | +| -------- | -------------- | ---------------- | --: | --: | --------- | --- | --: | +| `dp4cp1` | `[1,4]` | `[1,4]` | 4 | 1 | BF16 | THD | 1 | +| `dp2cp2` | `[1,4]` | `[1,4]` | 2 | 2 | BF16 | THD | 1 | +| `dp2cp1` | `[1,2]` | `[1,4]` | 2 | 1 | BF16 | THD | 1 | + +There is intentionally no precision argument or FP32 branch in the launcher. + +## Capture scope + +The hooks capture the forward used to calculate current policy log-probability, +selected by `before_log_prob`. The later P3O stats replay and gradient forward +consume the same Step-0 inputs and are deliberately not duplicated on disk. +Activation recomputation is disabled in every template. A run therefore has +exactly one dump for each non-dummy rank-local log-prob micro-batch. + +For each decoder layer `NNN`, the required stage set is: + +- `layer_NNN.block.input` and `.output`; +- `layer_NNN.self_attention.input` and `.output`; +- `layer_NNN.qkv_projection.input` and `.output`; +- `layer_NNN.attention_query`, `.attention_key`, and `.attention_value` at the + core-attention call boundary; +- `layer_NNN.attention_output` from core attention; +- `logits` from the model output. + +Hooks serialize tensors immediately after each stage to bound host memory. +The diagnostic sync and I/O overhead is intentional and must not be used for +performance measurements. + +## Artifact contract + +For a run directory ``, `p3o_forward_dump.py finalize-manifest` produces +this layout: + +```text +/ +├── command.sh +├── resolved_args.txt +├── stdout_stderr.log +├── exit_code.txt +├── manifest.json +└── dump/ + ├── runtime_rank.json + ├── manifest_rank.json + └── rank/micro/ + ├── metadata.json + ├── token_metadata.pt + ├── layer_000.block.input.pt + ├── ... + └── logits.pt +``` + +`R` is the actor global rank, not DP or CP rank; `M` is that rank's zero-based +micro-batch index within the captured log-prob pass. Every stage tensor is +stored losslessly on CPU with its original dtype. `metadata.json` maps each +stage to its file, shape, dtype, token axis, and finite/non-finite result. + +### Global sample and token key + +The topology-independent sample key is +`sha256(full_sample_token_ids.to(int64).contiguous().bytes)`, matching the 8/14 +oracle audit. The global tensor-row key is the pair +`(sample_sha256, global_zero_based_token_index)`. + +`token_metadata.pt` supplies the exact row mapping: + +- `sample_keys`: full-sequence SHA-256 values for this micro-batch; +- `full_token_ids`: the corresponding pre-CP token tensors; +- `local_sample_indices[i]`: index into `sample_keys`, or `-1` for trailing + pack padding; +- `local_token_indices[i]`: global position inside that sample, or `-1` for + trailing pack padding; +- `local_chunk_indices[i]`: global CP chunk number. For CP size `C`, rank `r` + owns chunks `r` and `2C-1-r`; CP1 uses chunk `0`; +- `local_real_mask[i]`: false for within-sample CP padding and trailing pack + padding. The comparator never treats padding as a real global token; +- `derived_position_ids`: the monotonic global positions that Relax/Megatron + synthesizes because the model call passes `position_ids=None`; +- `position_ids_argument`: the direct model argument (`None` for this frozen + command contract); +- `cu_seqlens_q` and `cu_seqlens_kv`: direct CPU copies from + `PackedSeqParams`. + +An identical full token sequence appearing twice would make a content-only +sample key ambiguous. Both the dumper contract and comparator treat that as a +hard harness failure rather than silently merging samples. + +### Per-capture metadata + +`metadata.json` is written only after the root model forward completes. It +contains: + +- `format_version`, `complete`, `phase`, and `phase_micro_batch_index`; +- global rank/world size plus DP/CP/TP/PP rank and world size; +- `qkv_format`, `micro_batch_size`, sample keys, total/response/max sequence + lengths, `max_seqlen_q`, and `max_seqlen_kv`; +- whether the direct position argument was `None`, the token-key contract, and + the `token_metadata.pt` path; +- all stage paths, shapes, dtypes, token axes, and finite flags. + +`runtime_rank.json` additionally records resolved precision, fixture path and +SHA-256, capture phase, and the two migration sources. `manifest_rank.json` +lists completed micro-batches for one rank. Top-level `manifest.json` requires +all expected ranks and capture metadata, checks every tensor file and finite +flag, and records fixture hashes and file counts. + +## Failure contract + +A cluster cell is failed if any of the following holds, even when Ray exits 0: + +- resolved precision is not BF16, layout is not THD, or MBS is not 1; +- fixture SHA-256 differs across ranks or from + `48538d165386dc94006613d857c022a7ba2e979bdc31bc617374eee2dc3c35b8`; +- a rank, non-dummy log-prob micro-batch, required stage, tensor file, + `token_metadata.pt`, or completed `metadata.json` is missing; +- a tensor contains NaN/Inf, lacks an unambiguous token axis, or its token-axis + length disagrees with the local ownership metadata; +- a sample hash is ambiguous, a global sample+token key is duplicated, CP ranks + do not reconstruct the same real-token key set as CP1, or token IDs/derived + positions disagree; +- `manifest.json.complete` is false; +- the comparator reports any contract error or a tensor exceeds its supplied + `atol`/`rtol`. Defaults are both `1e-6`; Batch 5 must report max-abs, + relative-L2, first divergent stage/token, and CP chunk-boundary statistics. + +The comparator exits 0 only for `FORWARD_MATCH`; `FORWARD_MISMATCH` exits 1 and +preserves `ROOT_CAUSE.json` and `ROOT_CAUSE.md`. Classification follows the +evidence location: input/key coverage failure → `INPUT_PACKING_BUG`, first +internal layer divergence → `ATTENTION_KERNEL_ORDER`, and logits-only +divergence after equal internal stages → `LOGPROB_EXTRACTION_BUG`. + +## Commands (Batch 4 only) + +Do not run these during Batch 3. On the designated 4×A100 cluster, invoke one +fresh run ID per cell; the common launcher refuses to overwrite a run directory. + +```bash +bash scripts/local_p3o_task40/run_forward_dump.sh dp4cp1 +bash scripts/local_p3o_task40/run_forward_dump.sh dp2cp2 +bash scripts/local_p3o_task40/run_forward_dump.sh dp2cp1 +``` + +The resolved run path is +`.../runs/step0_dump//seed_42//`. Compare completed runs with: + +```bash +python scripts/local_p3o_task40/compare_forward_dumps.py \ + --reference \ + --candidate \ + --output-dir /analysis +``` + +Batch 3 validation is limited to local imports, compilation, shell syntax, and +static command inspection. It does not submit Ray jobs or start Batch 4. diff --git a/scripts/local_p3o_task40/compare_forward_dumps.py b/scripts/local_p3o_task40/compare_forward_dumps.py new file mode 100755 index 000000000..759cbc39a --- /dev/null +++ b/scripts/local_p3o_task40/compare_forward_dumps.py @@ -0,0 +1,485 @@ +#!/usr/bin/env python3 + +"""Compare Task40 forward dumps by global sample and token key.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import sys +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Any + +import torch + + +STAGE_ORDER = { + "block.input": 0, + "self_attention.input": 1, + "qkv_projection.input": 2, + "qkv_projection.output": 3, + "attention_query": 4, + "attention_key": 5, + "attention_value": 6, + "attention_output": 7, + "self_attention.output": 8, + "block.output": 9, +} + + +@dataclass(frozen=True) +class Capture: + directory: Path + metadata: dict[str, Any] + + +@dataclass +class DumpIndex: + run_dir: Path + captures: list[Capture] + samples: dict[str, list[Capture]] + stages: set[str] + errors: list[str] + + +def _load_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text()) + + +def _find_dump_dir(run_dir: Path) -> Path: + manifest_path = run_dir / "manifest.json" + if manifest_path.is_file(): + manifest = _load_json(manifest_path) + return run_dir / manifest.get("dump_dir", "dump") + if (run_dir / "dump").is_dir(): + return run_dir / "dump" + if list(run_dir.glob("rank*/micro*/metadata.json")): + return run_dir + raise ValueError(f"cannot find a forward dump under {run_dir}") + + +def _load_index(run_dir: Path) -> DumpIndex: + dump_dir = _find_dump_dir(run_dir) + captures: list[Capture] = [] + samples: dict[str, list[Capture]] = {} + stages: set[str] = set() + errors: list[str] = [] + for metadata_path in sorted(dump_dir.glob("rank*/micro*/metadata.json")): + metadata = _load_json(metadata_path) + capture = Capture(metadata_path.parent, metadata) + captures.append(capture) + if not metadata.get("complete", False): + errors.append(f"incomplete capture: {metadata_path}") + if metadata.get("phase") != "log_prob": + errors.append(f"unexpected phase in {metadata_path}: {metadata.get('phase')}") + required_stages = set(metadata.get("required_stages", [])) + actual_stages = set(metadata.get("stages", {})) + if not required_stages or required_stages != actual_stages: + errors.append( + f"stage contract mismatch in {metadata_path}: " + f"missing={sorted(required_stages - actual_stages)}, " + f"unexpected={sorted(actual_stages - required_stages)}" + ) + token_metadata_path = capture.directory / metadata.get("token_metadata_path", "token_metadata.pt") + if not token_metadata_path.is_file(): + errors.append(f"missing token metadata: {token_metadata_path}") + else: + token_meta = _token_metadata(capture.directory) + cu_q = token_meta.get("cu_seqlens_q") + cu_kv = token_meta.get("cu_seqlens_kv") + if cu_q is None or cu_kv is None or not torch.equal(cu_q, cu_kv): + errors.append(f"missing or unequal cu_seqlens_q/kv: {metadata_path}") + elif cu_q.numel() < 2 or bool((cu_q[1:] < cu_q[:-1]).any()): + errors.append(f"non-monotonic cu_seqlens: {metadata_path}") + if token_meta.get("position_ids_argument") is not None: + errors.append(f"position_ids argument must be None for the frozen command: {metadata_path}") + if not torch.equal(token_meta["derived_position_ids"], token_meta["local_token_indices"]): + errors.append(f"derived position IDs disagree with token mapping: {metadata_path}") + for sample_key in metadata.get("sample_keys", []): + samples.setdefault(sample_key, []).append(capture) + stages.update(metadata.get("stages", {})) + for stage, stage_meta in metadata.get("stages", {}).items(): + if stage_meta.get("token_axis") is None: + errors.append(f"stage has no token axis: {metadata_path}:{stage}") + if not stage_meta.get("finite", False): + errors.append(f"non-finite stage: {metadata_path}:{stage}") + if not (capture.directory / stage_meta["path"]).is_file(): + errors.append(f"missing tensor: {metadata_path}:{stage}") + if not captures: + errors.append(f"no complete capture metadata under {dump_dir}") + return DumpIndex(run_dir.resolve(), captures, samples, stages, errors) + + +@lru_cache(maxsize=256) +def _token_metadata(directory: Path) -> dict[str, Any]: + return torch.load(directory / "token_metadata.pt", map_location="cpu", weights_only=False) + + +def _stage_sort_key(stage: str) -> tuple[int, int, str]: + match = re.match(r"layer_(\d+)\.(.+)", stage) + if not match: + return (10**9, 0 if stage == "logits" else 1, stage) + layer = int(match.group(1)) + suffix = match.group(2) + return (layer, STAGE_ORDER.get(suffix, 100), suffix) + + +def _rows_for_sample( + captures: list[Capture], sample_key: str, stage: str +) -> tuple[torch.Tensor, list[int], dict[int, int]]: + row_tensors: list[torch.Tensor] = [] + token_positions: list[int] = [] + chunks_by_position: dict[int, int] = {} + feature_shape: tuple[int, ...] | None = None + for capture in captures: + stage_meta = capture.metadata.get("stages", {}).get(stage) + if stage_meta is None: + continue + token_meta = _token_metadata(capture.directory) + sample_keys = token_meta["sample_keys"] + matching_sample_indices = [index for index, key in enumerate(sample_keys) if key == sample_key] + if len(matching_sample_indices) != 1: + raise ValueError( + f"{capture.directory}: sample key {sample_key} occurs {len(matching_sample_indices)} times; " + "content-hash key is ambiguous" + ) + sample_index = matching_sample_indices[0] + local_sample_indices = token_meta["local_sample_indices"] + local_token_indices = token_meta["local_token_indices"] + local_chunk_indices = token_meta["local_chunk_indices"] + local_real_mask = token_meta["local_real_mask"] + selected = ((local_sample_indices == sample_index) & local_real_mask).nonzero().flatten() + positions = local_token_indices[selected].to(torch.int64) + chunks = local_chunk_indices[selected].to(torch.int64) + + tensor = torch.load(capture.directory / stage_meta["path"], map_location="cpu", weights_only=False) + token_axis = int(stage_meta["token_axis"]) + tensor = tensor.movedim(token_axis, 0) + if tensor.shape[0] != local_token_indices.numel(): + raise ValueError( + f"{capture.directory}:{stage}: token axis has {tensor.shape[0]} rows, " + f"metadata has {local_token_indices.numel()}" + ) + current_feature_shape = tuple(int(size) for size in tensor.shape[1:]) + if feature_shape is None: + feature_shape = current_feature_shape + elif current_feature_shape != feature_shape: + raise ValueError(f"{capture.directory}:{stage}: feature shape {current_feature_shape} != {feature_shape}") + selected_rows = tensor.index_select(0, selected).reshape(selected.numel(), -1) + row_tensors.append(selected_rows) + for position, chunk in zip(positions.tolist(), chunks.tolist(), strict=True): + if position in chunks_by_position: + raise ValueError(f"duplicate global token key ({sample_key}, {position}) for stage {stage}") + chunks_by_position[position] = chunk + token_positions.append(position) + + if not row_tensors: + raise ValueError(f"sample {sample_key} has no tensor for stage {stage}") + rows = torch.cat(row_tensors, dim=0) + order = sorted(range(len(token_positions)), key=token_positions.__getitem__) + ordered_rows = rows.index_select(0, torch.tensor(order, dtype=torch.int64)) + ordered_positions = [token_positions[index] for index in order] + return ordered_rows, ordered_positions, chunks_by_position + + +def _compare_rows( + reference: torch.Tensor, + candidate: torch.Tensor, + positions: list[int], + candidate_chunks: dict[int, int], + *, + atol: float, + rtol: float, + max_chunk_elements: int, +) -> dict[str, Any]: + if reference.shape != candidate.shape: + raise ValueError(f"row shape mismatch: {tuple(reference.shape)} != {tuple(candidate.shape)}") + feature_size = int(reference.shape[1]) + rows_per_chunk = max(1, max_chunk_elements // max(1, feature_size)) + max_abs = 0.0 + diff_l2_sq = 0.0 + reference_l2_sq = 0.0 + first_bad_token: int | None = None + finite = True + by_cp_chunk: dict[int, float] = {} + near_boundary: dict[int, float] = {} + chunk_boundaries = sorted( + position + for previous, position in zip(positions, positions[1:]) + if candidate_chunks.get(previous) != candidate_chunks.get(position) + ) + + for start in range(0, reference.shape[0], rows_per_chunk): + stop = min(start + rows_per_chunk, reference.shape[0]) + left = reference[start:stop].float() + right = candidate[start:stop].float() + difference = left - right + finite = finite and bool(torch.isfinite(left).all() and torch.isfinite(right).all()) + absolute = difference.abs() + max_abs = max(max_abs, float(absolute.max()) if absolute.numel() else 0.0) + diff_l2_sq += float(torch.sum(difference.double() * difference.double())) + reference_l2_sq += float(torch.sum(left.double() * left.double())) + tolerance = atol + rtol * torch.maximum(left.abs(), right.abs()) + row_bad = (absolute > tolerance).reshape(stop - start, -1).any(dim=1) + row_max = absolute.reshape(stop - start, -1).amax(dim=1) + for offset, token_max in enumerate(row_max.tolist()): + position = positions[start + offset] + cp_chunk = candidate_chunks[position] + by_cp_chunk[cp_chunk] = max(by_cp_chunk.get(cp_chunk, 0.0), token_max) + distance = min((abs(position - boundary) for boundary in chunk_boundaries), default=10**9) + if distance <= 2: + near_boundary[distance] = max(near_boundary.get(distance, 0.0), token_max) + if first_bad_token is None and bool(row_bad.any()): + first_local = int(row_bad.nonzero()[0]) + first_bad_token = positions[start + first_local] + + relative_l2 = math.sqrt(diff_l2_sq) / max(math.sqrt(reference_l2_sq), 1e-30) + return { + "finite": finite, + "shape": list(reference.shape), + "token_count": len(positions), + "max_abs": max_abs, + "relative_l2": relative_l2, + "within_tolerance": finite and first_bad_token is None, + "first_bad_global_token_index": first_bad_token, + "max_abs_by_candidate_cp_chunk": {str(key): value for key, value in sorted(by_cp_chunk.items())}, + "max_abs_near_candidate_chunk_boundary": { + f"distance_{key}": value for key, value in sorted(near_boundary.items()) + }, + } + + +def _metadata_checks(reference: DumpIndex, candidate: DumpIndex) -> dict[str, Any]: + reference_samples = set(reference.samples) + candidate_samples = set(candidate.samples) + checks: dict[str, Any] = { + "sample_keys_equal": reference_samples == candidate_samples, + "reference_only_sample_keys": sorted(reference_samples - candidate_samples), + "candidate_only_sample_keys": sorted(candidate_samples - reference_samples), + "per_sample": {}, + } + for sample_key in sorted(reference_samples & candidate_samples): + reference_captures = reference.samples[sample_key] + candidate_captures = candidate.samples[sample_key] + reference_meta = _token_metadata(reference_captures[0].directory) + candidate_meta = _token_metadata(candidate_captures[0].directory) + reference_tokens = next( + tokens + for key, tokens in zip(reference_meta["sample_keys"], reference_meta["full_token_ids"], strict=True) + if key == sample_key + ) + candidate_tokens = next( + tokens + for key, tokens in zip(candidate_meta["sample_keys"], candidate_meta["full_token_ids"], strict=True) + if key == sample_key + ) + reference_positions = sorted( + { + int(position) + for capture in reference_captures + for key, position, real in zip( + ( + _token_metadata(capture.directory)["sample_keys"][int(index)] if int(index) >= 0 else None + for index in _token_metadata(capture.directory)["local_sample_indices"] + ), + _token_metadata(capture.directory)["local_token_indices"], + _token_metadata(capture.directory)["local_real_mask"], + strict=True, + ) + if key == sample_key and bool(real) + } + ) + candidate_positions = sorted( + { + int(position) + for capture in candidate_captures + for key, position, real in zip( + ( + _token_metadata(capture.directory)["sample_keys"][int(index)] if int(index) >= 0 else None + for index in _token_metadata(capture.directory)["local_sample_indices"] + ), + _token_metadata(capture.directory)["local_token_indices"], + _token_metadata(capture.directory)["local_real_mask"], + strict=True, + ) + if key == sample_key and bool(real) + } + ) + checks["per_sample"][sample_key] = { + "token_ids_equal": torch.equal(reference_tokens, candidate_tokens), + "global_token_positions_equal": reference_positions == candidate_positions, + "reference_token_count": len(reference_positions), + "candidate_token_count": len(candidate_positions), + } + checks["all_equal"] = checks["sample_keys_equal"] and all( + sample["token_ids_equal"] and sample["global_token_positions_equal"] + for sample in checks["per_sample"].values() + ) + return checks + + +def _classification(first_stage: str | None, metadata_equal: bool) -> str | None: + if not metadata_equal: + return "INPUT_PACKING_BUG" + if first_stage is None: + return None + if first_stage == "logits": + return "LOGPROB_EXTRACTION_BUG" + return "ATTENTION_KERNEL_ORDER" + + +def compare( + reference_dir: Path, + candidate_dir: Path, + *, + atol: float, + rtol: float, + max_chunk_elements: int, +) -> dict[str, Any]: + """Compare two runs under the global sample/token contract.""" + reference = _load_index(reference_dir) + candidate = _load_index(candidate_dir) + metadata = _metadata_checks(reference, candidate) + errors = reference.errors + candidate.errors + if reference.stages != candidate.stages: + errors.append( + f"stage sets differ: reference_only={sorted(reference.stages - candidate.stages)}, " + f"candidate_only={sorted(candidate.stages - reference.stages)}" + ) + + stage_reports: dict[str, Any] = {} + first_divergence: dict[str, Any] | None = None + common_samples = sorted(set(reference.samples) & set(candidate.samples)) + for stage in sorted(reference.stages & candidate.stages, key=_stage_sort_key): + sample_reports: dict[str, Any] = {} + for sample_key in common_samples: + try: + left, left_positions, _ = _rows_for_sample(reference.samples[sample_key], sample_key, stage) + right, right_positions, right_chunks = _rows_for_sample( + candidate.samples[sample_key], sample_key, stage + ) + if left_positions != right_positions: + raise ValueError( + f"global token coverage differs: reference={left_positions[:8]}... candidate={right_positions[:8]}..." + ) + result = _compare_rows( + left, + right, + left_positions, + right_chunks, + atol=atol, + rtol=rtol, + max_chunk_elements=max_chunk_elements, + ) + sample_reports[sample_key] = result + if not result["within_tolerance"] and first_divergence is None: + first_divergence = { + "stage": stage, + "sample_key": sample_key, + "global_token_index": result["first_bad_global_token_index"], + "max_abs": result["max_abs"], + "relative_l2": result["relative_l2"], + } + except (OSError, RuntimeError, ValueError) as exc: + message = f"{stage}:{sample_key}: {exc}" + errors.append(message) + sample_reports[sample_key] = {"error": str(exc), "within_tolerance": False} + if first_divergence is None: + first_divergence = {"stage": stage, "sample_key": sample_key, "error": str(exc)} + stage_reports[stage] = { + "within_tolerance": bool(sample_reports) + and all(sample.get("within_tolerance", False) for sample in sample_reports.values()), + "samples": sample_reports, + } + + passed = ( + metadata["all_equal"] + and not errors + and bool(stage_reports) + and all(stage["within_tolerance"] for stage in stage_reports.values()) + ) + first_stage = None if first_divergence is None else first_divergence.get("stage") + return { + "verdict": "FORWARD_MATCH" if passed else "FORWARD_MISMATCH", + "reference": str(reference.run_dir), + "candidate": str(candidate.run_dir), + "tolerances": {"atol": atol, "rtol": rtol}, + "metadata": metadata, + "errors": errors, + "stages": stage_reports, + "first_divergence": first_divergence, + "root_cause_classification": _classification(first_stage, metadata["all_equal"]), + } + + +def _render_markdown(report: dict[str, Any]) -> str: + lines = [ + "# Task40 CP forward comparison", + "", + f"**Verdict: `{report['verdict']}`**", + "", + f"- Reference: `{report['reference']}`", + f"- Candidate: `{report['candidate']}`", + f"- Metadata equal: `{report['metadata']['all_equal']}`", + f"- Root-cause classification: `{report['root_cause_classification']}`", + "", + "## First divergence", + "", + "```json", + json.dumps(report["first_divergence"], indent=2, sort_keys=True), + "```", + "", + "## Stage summary", + "", + "| Stage | Within tolerance | Worst max-abs | Worst rel-L2 |", + "| --- | ---: | ---: | ---: |", + ] + for stage, stage_report in report["stages"].items(): + numeric = [sample for sample in stage_report["samples"].values() if "max_abs" in sample] + worst_abs = max((sample["max_abs"] for sample in numeric), default=None) + worst_rel = max((sample["relative_l2"] for sample in numeric), default=None) + lines.append(f"| `{stage}` | `{stage_report['within_tolerance']}` | `{worst_abs}` | `{worst_rel}` |") + if report["errors"]: + lines.extend(["", "## Contract errors", "", "```json", json.dumps(report["errors"], indent=2), "```"]) + lines.append("") + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", type=Path, required=True, help="DP4CP1 or other CP1 run directory") + parser.add_argument("--candidate", type=Path, required=True, help="DP2CP2/DP2CP1 run directory") + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--atol", type=float, default=1e-6) + parser.add_argument("--rtol", type=float, default=1e-6) + parser.add_argument("--max-chunk-elements", type=int, default=1_000_000) + args = parser.parse_args() + if args.atol < 0 or args.rtol < 0: + parser.error("--atol and --rtol must be non-negative") + if args.max_chunk_elements <= 0: + parser.error("--max-chunk-elements must be positive") + + report = compare( + args.reference, + args.candidate, + atol=args.atol, + rtol=args.rtol, + max_chunk_elements=args.max_chunk_elements, + ) + args.output_dir.mkdir(parents=True, exist_ok=True) + (args.output_dir / "ROOT_CAUSE.json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + (args.output_dir / "ROOT_CAUSE.md").write_text(_render_markdown(report)) + sys.stdout.write( + json.dumps({"verdict": report["verdict"], "first_divergence": report["first_divergence"]}, sort_keys=True) + + "\n" + ) + if report["verdict"] != "FORWARD_MATCH": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/local_p3o_task40/p3o_forward_dump.py b/scripts/local_p3o_task40/p3o_forward_dump.py new file mode 100755 index 000000000..03997f2ee --- /dev/null +++ b/scripts/local_p3o_task40/p3o_forward_dump.py @@ -0,0 +1,666 @@ +#!/usr/bin/env python3 + +"""Task40 BF16 THD forward-dump hook and manifest finalizer. + +The module deliberately keeps PyTorch, Megatron, and Relax imports inside the +runtime entry points so that the command-line manifest finalizer and a plain +local import do not require the cluster environment. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import re +import sys +from collections import defaultdict, deque +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable + + +FORMAT_VERSION = 1 +CAPTURE_PHASE = "log_prob" +P0_ORACLE_SOURCE_SHA256 = "d53fdded5f601e0611231520166850dd49f05aad754e36f8adcef0333d31bd8c" +EXPECTED_FIXTURE_SHA256 = "48538d165386dc94006613d857c022a7ba2e979bdc31bc617374eee2dc3c35b8" + + +def _json_write(path: Path, payload: Any) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + temporary.replace(path) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _token_sha256(tokens: Any) -> str: + cpu_tokens = tokens.detach().to(device="cpu", dtype=_torch().int64).contiguous() + return hashlib.sha256(cpu_tokens.numpy().tobytes()).hexdigest() + + +def _torch() -> Any: + import torch + + return torch + + +def _first_tensor(value: Any) -> Any | None: + torch = _torch() + if torch.is_tensor(value): + return value + if isinstance(value, (tuple, list)): + for item in value: + tensor = _first_tensor(item) + if tensor is not None: + return tensor + if isinstance(value, dict): + for item in value.values(): + tensor = _first_tensor(item) + if tensor is not None: + return tensor + return None + + +def _argument_tensor(args: tuple[Any, ...], kwargs: dict[str, Any], names: tuple[str, ...], index: int) -> Any | None: + torch = _torch() + for name in names: + candidate = kwargs.get(name) + if torch.is_tensor(candidate): + return candidate + if index < len(args) and torch.is_tensor(args[index]): + return args[index] + return None + + +@dataclass +class BatchRecord: + phase: str + phase_micro_batch_index: int + sample_keys: list[str] + total_lengths: list[int] + response_lengths: list[int] + max_seq_lens: list[int] | None + full_token_ids: list[Any] + local_sample_indices: Any + local_token_indices: Any + local_chunk_indices: Any + local_real_mask: Any + derived_position_ids: Any + cu_seqlens_q: Any | None + cu_seqlens_kv: Any | None + max_seqlen_q: int | None + max_seqlen_kv: int | None + qkv_format: str + + +@dataclass +class Capture: + record: BatchRecord + directory: Path + stages: dict[str, dict[str, Any]] = field(default_factory=dict) + + +@dataclass +class DumpState: + args: Any | None = None + dump_dir: Path | None = None + rank: int = 0 + world_size: int = 1 + dp_rank: int = 0 + dp_world_size: int = 1 + cp_rank: int = 0 + cp_world_size: int = 1 + tp_rank: int = 0 + tp_world_size: int = 1 + pp_rank: int = 0 + pp_world_size: int = 1 + phase: str = "unclassified" + phase_counters: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + pending: deque[BatchRecord] = field(default_factory=deque) + active: Capture | None = None + hook_installed_model_ids: set[int] = field(default_factory=set) + capture_directories: list[str] = field(default_factory=list) + original_get_batch: Callable[..., Any] | None = None + original_cp_probe: Callable[..., Any] | None = None + required_stages: set[str] = field(default_factory=set) + + +_STATE = DumpState() + + +def _packed_value(packed: Any, name: str) -> Any | None: + value = getattr(packed, name, None) if packed is not None else None + if value is None: + return None + if hasattr(value, "detach"): + return value.detach().cpu() + return value + + +def _build_local_token_map(batch: dict[str, Any], cp_size: int, cp_rank: int) -> tuple[Any, Any, Any, Any, Any]: + """Reproduce ``slice_with_cp`` ownership without changing the batch.""" + torch = _torch() + full_tokens = batch["unconcat_tokens"] + local_sample_indices: list[int] = [] + local_token_indices: list[int] = [] + local_chunk_indices: list[int] = [] + local_real_mask: list[bool] = [] + + for sample_index, tokens in enumerate(full_tokens): + total_length = int(tokens.shape[0]) + if cp_size == 1: + positions = range(total_length) + chunks = [0] * total_length + else: + chunk_size = math.ceil(total_length / (2 * cp_size)) + chunk_ids = (cp_rank, 2 * cp_size - cp_rank - 1) + positions = [ + position + for chunk_id in chunk_ids + for position in range(chunk_id * chunk_size, (chunk_id + 1) * chunk_size) + ] + chunks = [chunk_id for chunk_id in chunk_ids for _ in range(chunk_size)] + for position, chunk_id in zip(positions, chunks, strict=True): + local_sample_indices.append(sample_index) + local_token_indices.append(position) + local_chunk_indices.append(chunk_id) + local_real_mask.append(position < total_length) + + local_tokens = batch["tokens"] + local_length = int(local_tokens.numel()) + trailing_padding = local_length - len(local_token_indices) + if trailing_padding < 0: + raise ValueError( + f"forward dump token map exceeds local tensor length: mapped={len(local_token_indices)}, local={local_length}" + ) + local_sample_indices.extend([-1] * trailing_padding) + local_token_indices.extend([-1] * trailing_padding) + local_chunk_indices.extend([-1] * trailing_padding) + local_real_mask.extend([False] * trailing_padding) + + sample_tensor = torch.tensor(local_sample_indices, dtype=torch.int32) + token_tensor = torch.tensor(local_token_indices, dtype=torch.int64) + chunk_tensor = torch.tensor(local_chunk_indices, dtype=torch.int32) + real_tensor = torch.tensor(local_real_mask, dtype=torch.bool) + return sample_tensor, token_tensor, chunk_tensor, real_tensor, token_tensor.clone() + + +def _record_batch(batch: dict[str, Any]) -> BatchRecord: + phase = _STATE.phase + micro_batch_index = _STATE.phase_counters[phase] + _STATE.phase_counters[phase] += 1 + full_tokens = [tokens.detach().cpu() for tokens in batch["unconcat_tokens"]] + sample_keys = [_token_sha256(tokens) for tokens in full_tokens] + sample_indices, token_indices, chunk_indices, real_mask, positions = _build_local_token_map( + batch, _STATE.cp_world_size, _STATE.cp_rank + ) + packed = batch.get("packed_seq_params") + raw_max_seq_lens = batch.get("max_seq_lens") + max_seq_lens = None if raw_max_seq_lens is None else [int(value) for value in raw_max_seq_lens] + return BatchRecord( + phase=phase, + phase_micro_batch_index=micro_batch_index, + sample_keys=sample_keys, + total_lengths=[int(value) for value in batch["total_lengths"]], + response_lengths=[int(value) for value in batch["response_lengths"]], + max_seq_lens=max_seq_lens, + full_token_ids=full_tokens, + local_sample_indices=sample_indices, + local_token_indices=token_indices, + local_chunk_indices=chunk_indices, + local_real_mask=real_mask, + derived_position_ids=positions, + cu_seqlens_q=_packed_value(packed, "cu_seqlens_q"), + cu_seqlens_kv=_packed_value(packed, "cu_seqlens_kv"), + max_seqlen_q=getattr(packed, "max_seqlen_q", None) if packed is not None else None, + max_seqlen_kv=getattr(packed, "max_seqlen_kv", None) if packed is not None else None, + qkv_format=getattr(packed, "qkv_format", getattr(_STATE.args, "qkv_format", "thd")), + ) + + +def _observed_get_batch(*args: Any, **kwargs: Any) -> Any: + if _STATE.original_get_batch is None: + raise RuntimeError("forward dump get_batch wrapper installed without its original callable") + batch = _STATE.original_get_batch(*args, **kwargs) + if _STATE.phase == CAPTURE_PHASE and not batch.get("__is_dummy__", False): + _STATE.pending.append(_record_batch(batch)) + return batch + + +def _sanitize_stage(stage: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]", "_", stage) + + +def _infer_token_axis(tensor: Any, local_length: int, preferred: int | None) -> int | None: + if preferred is not None and tensor.dim() > preferred and int(tensor.shape[preferred]) == local_length: + return preferred + matches = [axis for axis, size in enumerate(tensor.shape) if int(size) == local_length] + return matches[0] if len(matches) == 1 else None + + +def _store_stage(stage: str, value: Any, preferred_token_axis: int | None = 0) -> None: + capture = _STATE.active + if capture is None: + return + tensor = _first_tensor(value) + if tensor is None: + return + torch = _torch() + cpu_tensor = tensor.detach().cpu().contiguous() + local_length = int(capture.record.local_token_indices.numel()) + token_axis = _infer_token_axis(cpu_tensor, local_length, preferred_token_axis) + filename = _sanitize_stage(stage) + ".pt" + torch.save(cpu_tensor, capture.directory / filename) + finite = bool(torch.isfinite(cpu_tensor).all()) if cpu_tensor.is_floating_point() else True + capture.stages[stage] = { + "dtype": str(cpu_tensor.dtype), + "finite": finite, + "path": filename, + "shape": list(cpu_tensor.shape), + "token_axis": token_axis, + } + + +def _write_rank_manifest() -> None: + if _STATE.dump_dir is None: + return + _json_write( + _STATE.dump_dir / f"manifest_rank{_STATE.rank}.json", + { + "format_version": FORMAT_VERSION, + "rank": _STATE.rank, + "world_size": _STATE.world_size, + "captures": _STATE.capture_directories, + }, + ) + + +def _root_pre_hook(module: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None: + del module, args + if _STATE.phase != CAPTURE_PHASE: + return + if not _STATE.pending: + raise RuntimeError("forward dump saw a log-prob model forward without matching get_batch metadata") + if _STATE.active is not None: + raise RuntimeError("forward dump does not support nested model forwards") + record = _STATE.pending.popleft() + if _STATE.dump_dir is None: + raise RuntimeError("forward dump directory was not configured") + directory = _STATE.dump_dir / f"rank{_STATE.rank:05d}" / f"micro{record.phase_micro_batch_index:05d}" + directory.mkdir(parents=True, exist_ok=False) + _STATE.active = Capture(record=record, directory=directory) + + position_ids = kwargs.get("position_ids") + token_metadata = { + "sample_keys": record.sample_keys, + "full_token_ids": record.full_token_ids, + "local_sample_indices": record.local_sample_indices, + "local_token_indices": record.local_token_indices, + "local_chunk_indices": record.local_chunk_indices, + "local_real_mask": record.local_real_mask, + "derived_position_ids": record.derived_position_ids, + "position_ids_argument": None if position_ids is None else position_ids.detach().cpu(), + "cu_seqlens_q": record.cu_seqlens_q, + "cu_seqlens_kv": record.cu_seqlens_kv, + } + _torch().save(token_metadata, directory / "token_metadata.pt") + + +def _root_post_hook(module: Any, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any) -> None: + del module, args, kwargs + capture = _STATE.active + if capture is None: + return + _store_stage("logits", output, preferred_token_axis=1) + record = capture.record + metadata = { + "format_version": FORMAT_VERSION, + "complete": True, + "phase": record.phase, + "phase_micro_batch_index": record.phase_micro_batch_index, + "rank": _STATE.rank, + "world_size": _STATE.world_size, + "dp_rank": _STATE.dp_rank, + "dp_world_size": _STATE.dp_world_size, + "cp_rank": _STATE.cp_rank, + "cp_world_size": _STATE.cp_world_size, + "tp_rank": _STATE.tp_rank, + "tp_world_size": _STATE.tp_world_size, + "pp_rank": _STATE.pp_rank, + "pp_world_size": _STATE.pp_world_size, + "qkv_format": record.qkv_format, + "micro_batch_size": int(getattr(_STATE.args, "micro_batch_size", 0)), + "sample_keys": record.sample_keys, + "total_lengths": record.total_lengths, + "response_lengths": record.response_lengths, + "max_seq_lens": record.max_seq_lens, + "max_seqlen_q": None if record.max_seqlen_q is None else int(record.max_seqlen_q), + "max_seqlen_kv": None if record.max_seqlen_kv is None else int(record.max_seqlen_kv), + "position_ids_argument_was_none": True, + "token_key_contract": "(sample_sha256_of_full_int64_token_ids, global_zero_based_token_index)", + "token_metadata_path": "token_metadata.pt", + "required_stages": sorted(_STATE.required_stages, key=_stage_sort_key), + "stages": capture.stages, + } + token_metadata = _torch().load(capture.directory / "token_metadata.pt", map_location="cpu", weights_only=False) + metadata["position_ids_argument_was_none"] = token_metadata["position_ids_argument"] is None + _json_write(capture.directory / "metadata.json", metadata) + relative = str(capture.directory.relative_to(_STATE.dump_dir)) + _STATE.capture_directories.append(relative) + _STATE.active = None + _write_rank_manifest() + + +def _register_pre_post(module: Any, stage_prefix: str, preferred_axis: int | None = 0) -> None: + def pre_hook(hook_module: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None: + del hook_module + tensor = _argument_tensor(args, kwargs, ("hidden_states",), 0) + _store_stage(f"{stage_prefix}.input", tensor, preferred_axis) + + def post_hook(hook_module: Any, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any) -> None: + del hook_module, args, kwargs + _store_stage(f"{stage_prefix}.output", output, preferred_axis) + + module.register_forward_pre_hook(pre_hook, with_kwargs=True) + module.register_forward_hook(post_hook, with_kwargs=True) + + +def _register_core_attention(module: Any, layer_prefix: str) -> None: + def pre_hook(hook_module: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None: + del hook_module + names = ( + (("query", "q", "query_layer"), 0, "query"), + (("key", "k", "key_layer"), 1, "key"), + (("value", "v", "value_layer"), 2, "value"), + ) + for aliases, index, label in names: + _store_stage( + f"{layer_prefix}.attention_{label}", + _argument_tensor(args, kwargs, aliases, index), + preferred_token_axis=0, + ) + + def post_hook(hook_module: Any, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any) -> None: + del hook_module, args, kwargs + _store_stage(f"{layer_prefix}.attention_output", output, preferred_token_axis=0) + + module.register_forward_pre_hook(pre_hook, with_kwargs=True) + module.register_forward_hook(post_hook, with_kwargs=True) + + +def _layer_number(name: str) -> int | None: + match = re.search(r"(?:^|\.)layers\.(\d+)$", name) + return int(match.group(1)) if match else None + + +def _stage_sort_key(stage: str) -> tuple[int, str]: + match = re.match(r"layer_(\d+)\.(.+)", stage) + return (int(match.group(1)), match.group(2)) if match else (10**9, stage) + + +def _install_forward_dump(model: Any) -> None: + model_id = id(model) + if model_id in _STATE.hook_installed_model_ids: + return + _STATE.hook_installed_model_ids.add(model_id) + + layer_names: dict[str, int] = {} + modules = list(model.named_modules()) + for name, module in modules: + layer = _layer_number(name) + if layer is not None: + layer_names[name] = layer + _register_pre_post(module, f"layer_{layer:03d}.block", preferred_axis=0) + prefix = f"layer_{layer:03d}" + _STATE.required_stages.update( + { + f"{prefix}.block.input", + f"{prefix}.block.output", + f"{prefix}.self_attention.input", + f"{prefix}.self_attention.output", + f"{prefix}.qkv_projection.input", + f"{prefix}.qkv_projection.output", + f"{prefix}.attention_query", + f"{prefix}.attention_key", + f"{prefix}.attention_value", + f"{prefix}.attention_output", + } + ) + + for name, module in modules: + owner = next((layer for prefix, layer in layer_names.items() if name.startswith(prefix + ".")), None) + if owner is None: + continue + layer_prefix = f"layer_{owner:03d}" + if name.endswith(".self_attention.linear_qkv"): + _register_pre_post(module, f"{layer_prefix}.qkv_projection", preferred_axis=0) + elif name.endswith(".self_attention.core_attention"): + _register_core_attention(module, layer_prefix) + elif name.endswith(".self_attention"): + _register_pre_post(module, f"{layer_prefix}.self_attention", preferred_axis=0) + + model.register_forward_pre_hook(_root_pre_hook, with_kwargs=True) + model.register_forward_hook(_root_post_hook, with_kwargs=True) + _STATE.required_stages.add("logits") + + +def _runtime_metadata(args: Any) -> dict[str, Any]: + fixture = Path(str(getattr(args, "load_debug_rollout_data", ""))) + fixture_sha = _sha256(fixture) if fixture.is_file() else None + return { + "format_version": FORMAT_VERSION, + "rank": _STATE.rank, + "world_size": _STATE.world_size, + "dp_rank": _STATE.dp_rank, + "dp_world_size": _STATE.dp_world_size, + "cp_rank": _STATE.cp_rank, + "cp_world_size": _STATE.cp_world_size, + "tp_rank": _STATE.tp_rank, + "tp_world_size": _STATE.tp_world_size, + "pp_rank": _STATE.pp_rank, + "pp_world_size": _STATE.pp_world_size, + "bf16": bool(getattr(args, "bf16", False)), + "fp16": bool(getattr(args, "fp16", False)), + "params_dtype": str(getattr(args, "params_dtype", None)), + "qkv_format": str(getattr(args, "qkv_format", None)), + "micro_batch_size": int(getattr(args, "micro_batch_size", 0)), + "global_batch_size": int(getattr(args, "global_batch_size", 0)), + "fixture": str(fixture), + "fixture_sha256": fixture_sha, + "expected_fixture_sha256": EXPECTED_FIXTURE_SHA256, + "capture_phase": CAPTURE_PHASE, + "p0_oracle_source_sha256": P0_ORACLE_SOURCE_SHA256, + "model_provider_cp_probe": "relax.backends.megatron.model_provider._install_cp_probe", + } + + +def configure(args: Any) -> None: + """Install the dump hooks through ``--custom-megatron-init-path``.""" + import torch.distributed as dist + from megatron.core import mpu + + from relax.backends.megatron import model as model_backend + from relax.backends.megatron import model_provider + + if _STATE.args is not None: + raise RuntimeError("Task40 forward dump configure() was called more than once in one process") + if str(getattr(args, "qkv_format", "thd")) != "thd": + raise ValueError("Task40 Batch 3 forward dump supports only THD") + if int(getattr(args, "micro_batch_size", 0)) != 1: + raise ValueError("Task40 Batch 3 forward dump supports only micro-batch size 1") + if not bool(getattr(args, "bf16", False)) or bool(getattr(args, "fp16", False)): + raise ValueError("Task40 Batch 3 forward dump requires BF16 and forbids FP16/FP32 templates") + fixture = Path(str(getattr(args, "load_debug_rollout_data", ""))) + if not fixture.is_file(): + raise ValueError(f"Task40 Step-0 fixture does not exist: {fixture}") + fixture_sha = _sha256(fixture) + if fixture_sha != EXPECTED_FIXTURE_SHA256: + raise ValueError( + f"Task40 Step-0 fixture SHA-256 mismatch: got {fixture_sha}, expected {EXPECTED_FIXTURE_SHA256}" + ) + + _STATE.args = args + _STATE.rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + _STATE.world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 + _STATE.dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False) + _STATE.dp_world_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + _STATE.cp_rank = mpu.get_context_parallel_rank() + _STATE.cp_world_size = mpu.get_context_parallel_world_size() + _STATE.tp_rank = mpu.get_tensor_model_parallel_rank() + _STATE.tp_world_size = mpu.get_tensor_model_parallel_world_size() + _STATE.pp_rank = mpu.get_pipeline_model_parallel_rank() + _STATE.pp_world_size = mpu.get_pipeline_model_parallel_world_size() + _STATE.dump_dir = Path(args.dump_details).parent / "dump" + _STATE.dump_dir.mkdir(parents=True, exist_ok=True) + (_STATE.dump_dir / f"rank{_STATE.rank:05d}").mkdir(exist_ok=True) + _json_write(_STATE.dump_dir / f"runtime_rank{_STATE.rank}.json", _runtime_metadata(args)) + + _STATE.original_get_batch = model_backend.get_batch + model_backend.get_batch = _observed_get_batch + + _STATE.original_cp_probe = model_provider._install_cp_probe + + def install_cp_probe_and_dump(model: Any) -> None: + if _STATE.original_cp_probe is None: + raise RuntimeError("forward dump lost the original CP probe") + _STATE.original_cp_probe(model) + _install_forward_dump(model) + + model_provider._install_cp_probe = install_cp_probe_and_dump + + +def before_log_prob(args: Any, model: Any, store_prefix: str) -> None: + """Select the one production log-prob pass captured by this harness.""" + del args, model + _STATE.phase = CAPTURE_PHASE if store_prefix == "" else f"ignored_{store_prefix}log_prob" + + +def before_train_step( + args: Any, + rollout_id: int, + step_id: int, + model: Any, + optimizer: Any, + opt_param_scheduler: Any, +) -> None: + """Stop capture before the replayed P3O stats/train forwards.""" + del args, rollout_id, step_id, model, optimizer, opt_param_scheduler + _STATE.phase = "train" + + +def finalize_manifest(run_dir: Path) -> dict[str, Any]: + """Validate rank-local captures and write the run-level manifest.""" + dump_dir = run_dir / "dump" + runtimes = [json.loads(path.read_text()) for path in sorted(dump_dir.glob("runtime_rank*.json"))] + rank_manifests = [json.loads(path.read_text()) for path in sorted(dump_dir.glob("manifest_rank*.json"))] + errors: list[str] = [] + if not runtimes: + errors.append("no runtime_rank*.json files") + expected_world_sizes = sorted({int(runtime["world_size"]) for runtime in runtimes}) + if len(expected_world_sizes) != 1: + errors.append(f"inconsistent world sizes: {expected_world_sizes}") + expected_ranks = expected_world_sizes[0] if len(expected_world_sizes) == 1 else 0 + runtime_ranks = {int(runtime["rank"]) for runtime in runtimes} + manifest_ranks = {int(manifest["rank"]) for manifest in rank_manifests} + if runtime_ranks != set(range(expected_ranks)): + errors.append(f"runtime ranks={sorted(runtime_ranks)}, expected={list(range(expected_ranks))}") + if manifest_ranks != set(range(expected_ranks)): + errors.append(f"manifest ranks={sorted(manifest_ranks)}, expected={list(range(expected_ranks))}") + + runtime_by_rank = {int(runtime["rank"]): runtime for runtime in runtimes} + for rank_manifest in rank_manifests: + rank = int(rank_manifest["rank"]) + runtime = runtime_by_rank.get(rank) + if runtime is None: + continue + denominator = int(runtime["dp_world_size"]) * int(runtime["micro_batch_size"]) + global_batch_size = int(runtime["global_batch_size"]) + if denominator <= 0 or global_batch_size % denominator != 0: + errors.append( + f"rank {rank} cannot derive micro-batch count from " + f"global_batch_size={global_batch_size}, dp_world_size={runtime['dp_world_size']}, " + f"micro_batch_size={runtime['micro_batch_size']}" + ) + continue + expected_micro_batches = global_batch_size // denominator + actual_micro_batches = len(rank_manifest.get("captures", [])) + if actual_micro_batches != expected_micro_batches: + errors.append(f"rank {rank} capture count={actual_micro_batches}, expected={expected_micro_batches}") + + captures: list[str] = [] + for manifest in rank_manifests: + captures.extend(str(path) for path in manifest.get("captures", [])) + stage_file_count = 0 + nonfinite_stages: list[str] = [] + for relative in captures: + metadata_path = dump_dir / relative / "metadata.json" + if not metadata_path.is_file(): + errors.append(f"missing metadata: {relative}") + continue + metadata = json.loads(metadata_path.read_text()) + if not metadata.get("complete", False): + errors.append(f"incomplete capture: {relative}") + required_stages = set(metadata.get("required_stages", [])) + actual_stages = set(metadata.get("stages", {})) + if not required_stages: + errors.append(f"capture declares no required stages: {relative}") + if required_stages != actual_stages: + errors.append( + f"stage contract mismatch for {relative}: " + f"missing={sorted(required_stages - actual_stages)}, unexpected={sorted(actual_stages - required_stages)}" + ) + for stage, stage_meta in metadata.get("stages", {}).items(): + stage_file_count += 1 + if not (dump_dir / relative / stage_meta["path"]).is_file(): + errors.append(f"missing tensor: {relative}/{stage_meta['path']}") + if not stage_meta.get("finite", False): + nonfinite_stages.append(f"{relative}:{stage}") + if nonfinite_stages: + errors.append(f"non-finite stages: {nonfinite_stages[:8]}") + + fixture_hashes = sorted({str(runtime.get("fixture_sha256")) for runtime in runtimes}) + if fixture_hashes != [EXPECTED_FIXTURE_SHA256]: + errors.append(f"fixture SHA-256 values={fixture_hashes}, expected={[EXPECTED_FIXTURE_SHA256]}") + manifest = { + "format_version": FORMAT_VERSION, + "complete": not errors, + "errors": errors, + "run_dir": str(run_dir.resolve()), + "dump_dir": "dump", + "runtime_files": [str(path.relative_to(run_dir)) for path in sorted(dump_dir.glob("runtime_rank*.json"))], + "rank_manifest_files": [ + str(path.relative_to(run_dir)) for path in sorted(dump_dir.glob("manifest_rank*.json")) + ], + "capture_count": len(captures), + "stage_file_count": stage_file_count, + "fixture_sha256_values": fixture_hashes, + "captures": sorted(captures), + } + _json_write(run_dir / "manifest.json", manifest) + return manifest + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + finalize_parser = subparsers.add_parser("finalize-manifest", help="validate rank dumps and write manifest.json") + finalize_parser.add_argument("--run-dir", type=Path, required=True) + args = parser.parse_args() + manifest = finalize_manifest(args.run_dir) + sys.stdout.write( + json.dumps({"complete": manifest["complete"], "errors": manifest["errors"]}, sort_keys=True) + "\n" + ) + if not manifest["complete"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/local_p3o_task40/run_forward_dump.sh b/scripts/local_p3o_task40/run_forward_dump.sh new file mode 100755 index 000000000..20ae95e2f --- /dev/null +++ b/scripts/local_p3o_task40/run_forward_dump.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash + +# Task40 Batch 3 command package: BF16 + THD + MBS1 only. + +set -euo pipefail + +if [[ "${1:-}" == "--in-container" ]]; then + TOPOLOGY="${2:?usage: run_forward_dump.sh --in-container }" + case "${TOPOLOGY}" in + dp4cp1) ACTOR_WORLD_SIZE=4; CONTEXT_PARALLEL_SIZE=1 ;; + dp2cp2) ACTOR_WORLD_SIZE=4; CONTEXT_PARALLEL_SIZE=2 ;; + dp2cp1) ACTOR_WORLD_SIZE=2; CONTEXT_PARALLEL_SIZE=1 ;; + *) echo "unsupported topology: ${TOPOLOGY}" >&2; exit 2 ;; + esac + + : "${P3O_STEP0_FIXTURE:?P3O_STEP0_FIXTURE must be set}" + export P3O_ALGORITHM=p3o + export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + + SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" + REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." >/dev/null 2>&1 && pwd)" + source "${REPO_ROOT}/examples/algorithms/p3o/common_a100x4.sh" + + eval "$(declare -f P3O_build_args | sed '1s/P3O_build_args/P3O_build_args_base/')" + + replace_train_arg() { + local flag="$1" + local value="$2" + local index + for ((index = 0; index < ${#P3O_TRAIN_ARGS[@]}; index++)); do + if [[ "${P3O_TRAIN_ARGS[index]}" == "${flag}" ]]; then + P3O_TRAIN_ARGS[index + 1]="${value}" + return 0 + fi + done + echo "required argument not found: ${flag}" >&2 + return 3 + } + + P3O_build_args() { + P3O_build_args_base + replace_train_arg --resource "{\"actor\":[1,${ACTOR_WORLD_SIZE}],\"rollout\":[1,4]}" + replace_train_arg --context-parallel-size "${CONTEXT_PARALLEL_SIZE}" + replace_train_arg --lr 1e-6 + P3O_TRAIN_ARGS+=( + --load-debug-rollout-data "${P3O_STEP0_FIXTURE}" + --custom-megatron-init-path scripts.local_p3o_task40.p3o_forward_dump.configure + --custom-megatron-before-log-prob-hook-path scripts.local_p3o_task40.p3o_forward_dump.before_log_prob + --custom-megatron-before-train-step-hook-path scripts.local_p3o_task40.p3o_forward_dump.before_train_step + --dump-details "${P3O_OUTPUT_ROOT}/${P3O_CONFIG_NAME}/seed_${P3O_SEED}/${P3O_RUN_ID}/debug" + ) + } + + P3O_CONFIG_NAME="${TOPOLOGY}" + set +e + P3O_run + TRAIN_EXIT_CODE=$? + set -e + FINALIZE_EXIT_CODE=0 + if [[ -n "${P3O_RUN_DIR:-}" && -d "${P3O_RUN_DIR}" ]]; then + python3 scripts/local_p3o_task40/p3o_forward_dump.py finalize-manifest \ + --run-dir "${P3O_RUN_DIR}" || FINALIZE_EXIT_CODE=$? + else + FINALIZE_EXIT_CODE=1 + fi + if [[ "${TRAIN_EXIT_CODE}" -ne 0 ]]; then + exit "${TRAIN_EXIT_CODE}" + fi + exit "${FINALIZE_EXIT_CODE}" +fi + +TOPOLOGY="${1:?usage: run_forward_dump.sh }" +RUN_ID="${2:?usage: run_forward_dump.sh }" +case "${TOPOLOGY}" in + dp4cp1|dp2cp2|dp2cp1) ;; + *) echo "unsupported topology: ${TOPOLOGY}" >&2; exit 2 ;; +esac + +INFRA=/lustre/home/sztu_camdt_zhanghua/jimaomo/infra +REPO="${INFRA}/Relax" +IMAGE="${INFRA}/images/relaxrl-dev-20260715-8325919e.sif" +CAMPAIGN="${INFRA}/Output/task40/task40_cp_forward_diag_20260817_6c7a3d2" +FIXTURE="${INFRA}/Output/task40/task40_p0_cluster_20260814_6c7a3d2/fixtures/step0_rollout.pt" +LOG_DIR="${CAMPAIGN}/logs" +mkdir -p "${LOG_DIR}" + +set +e +apptainer exec --nv --bind /lustre:/lustre "${IMAGE}" env \ + P3O_MODE=smoke \ + P3O_MODEL_CONFIG="${REPO}/scripts/local_p3o_task40/qwen2p5_1p5b.sh" \ + P3O_MODEL_ROTARY_BASE=1000000 \ + P3O_MODEL_DIR="${INFRA}/Qwen2.5-1.5B-Instruct" \ + P3O_TRAIN_DATA="${INFRA}/gsm8k/main/train_clean.parquet" \ + P3O_INPUT_KEY=question \ + P3O_LABEL_KEY=answer \ + P3O_RM_TYPE=openr1mm \ + P3O_OUTPUT_ROOT="${CAMPAIGN}/runs/step0_dump" \ + P3O_MEGATRON_DIR=/root/Megatron-LM \ + P3O_RAY_DASHBOARD=http://127.0.0.1:8265 \ + P3O_NUM_ROLLOUT=1 \ + P3O_ROLLOUT_BATCH_SIZE=4 \ + P3O_N_SAMPLES=16 \ + P3O_GLOBAL_BATCH_SIZE=64 \ + P3O_MICRO_BATCH_SIZE=1 \ + P3O_MAX_RESPONSE_LEN=4096 \ + P3O_ESS_SCOPE=step \ + P3O_KL_MODE=proxy_safe \ + P3O_SEED=42 \ + P3O_PIPELINE_MODEL_PARALLEL_SIZE=1 \ + P3O_ACTIVATION_RECOMPUTE=0 \ + P3O_LOG_PROBS_CHUNK_SIZE=1024 \ + P3O_ROLLOUT_SHUFFLE=0 \ + P3O_DETERMINISTIC_INFERENCE=1 \ + P3O_CLEAR_RUNTIME_PROXIES=1 \ + P3O_STEP0_FIXTURE="${FIXTURE}" \ + P3O_RUN_ID="${RUN_ID}" \ + bash -lc "cd '${REPO}' && bash scripts/local_p3o_task40/run_forward_dump.sh --in-container '${TOPOLOGY}'" \ + 2>&1 | tee "${LOG_DIR}/step0_dump_${TOPOLOGY}_${RUN_ID}.launcher.log" +EXIT_CODE=${PIPESTATUS[0]} +set -e +printf '%s\n' "${EXIT_CODE}" >"${LOG_DIR}/step0_dump_${TOPOLOGY}_${RUN_ID}.launcher_exit_code.txt" +exit "${EXIT_CODE}" From 65d3e86041562b32757c35586f726aac2df8b73a Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 15:59:22 +0800 Subject: [PATCH 14/30] fix(p3o): repair CP forward dump harness --- scripts/local_p3o_task40/README.md | 19 +++++-- scripts/local_p3o_task40/p3o_forward_dump.py | 35 ++++++++++-- .../local_p3o_task40/test_p3o_forward_dump.py | 57 +++++++++++++++++++ 3 files changed, 100 insertions(+), 11 deletions(-) create mode 100644 tests/scripts/local_p3o_task40/test_p3o_forward_dump.py diff --git a/scripts/local_p3o_task40/README.md b/scripts/local_p3o_task40/README.md index 67bea5745..48935ddc4 100644 --- a/scripts/local_p3o_task40/README.md +++ b/scripts/local_p3o_task40/README.md @@ -54,11 +54,20 @@ There is intentionally no precision argument or FP32 branch in the launcher. ## Capture scope -The hooks capture the forward used to calculate current policy log-probability, -selected by `before_log_prob`. The later P3O stats replay and gradient forward -consume the same Step-0 inputs and are deliberately not duplicated on disk. -Activation recomputation is disabled in every template. A run therefore has -exactly one dump for each non-dummy rank-local log-prob micro-batch. +The frozen commands pass `--use-rollout-logprobs`, so the actor intentionally +skips a separate current-policy log-prob forward. The hooks therefore select +the first production P3O stats forward through `before_train_step`; for a +configuration that does execute actor log-prob forward, `before_log_prob` +selects that earlier equivalent path. Capture stops after exactly +`global_batch_size / (DP * MBS)` completed rank-local micro-batches, before the +gradient forward can be duplicated on disk. Activation recomputation is +disabled in every template. A run therefore has exactly one dump for each +non-dummy rank-local micro-batch. + +The hook observes both module-local aliases of `get_batch`: the ordinary +forward path in `model.py` and the P3O sufficient-statistics path in +`p3o_step.py`. This is required because each module imports the callable by +value; replacing only one alias does not observe the other path. For each decoder layer `NNN`, the required stage set is: diff --git a/scripts/local_p3o_task40/p3o_forward_dump.py b/scripts/local_p3o_task40/p3o_forward_dump.py index 03997f2ee..b0972dd0d 100755 --- a/scripts/local_p3o_task40/p3o_forward_dump.py +++ b/scripts/local_p3o_task40/p3o_forward_dump.py @@ -131,6 +131,7 @@ class DumpState: original_get_batch: Callable[..., Any] | None = None original_cp_probe: Callable[..., Any] | None = None required_stages: set[str] = field(default_factory=set) + capture_target_count: int = 0 _STATE = DumpState() @@ -235,6 +236,15 @@ def _observed_get_batch(*args: Any, **kwargs: Any) -> Any: return batch +def _install_get_batch_observers(model_backend: Any, p3o_step_backend: Any) -> None: + original = model_backend.get_batch + if original is not p3o_step_backend.get_batch: + raise RuntimeError("model.py and p3o_step.py do not share the expected original get_batch callable") + _STATE.original_get_batch = original + model_backend.get_batch = _observed_get_batch + p3o_step_backend.get_batch = _observed_get_batch + + def _sanitize_stage(stage: str) -> str: return re.sub(r"[^A-Za-z0-9_.-]", "_", stage) @@ -357,6 +367,8 @@ def _root_post_hook(module: Any, args: tuple[Any, ...], kwargs: dict[str, Any], _STATE.capture_directories.append(relative) _STATE.active = None _write_rank_manifest() + if len(_STATE.capture_directories) >= _STATE.capture_target_count: + _STATE.phase = "train_after_capture" def _register_pre_post(module: Any, stage_prefix: str, preferred_axis: int | None = 0) -> None: @@ -488,7 +500,7 @@ def configure(args: Any) -> None: from megatron.core import mpu from relax.backends.megatron import model as model_backend - from relax.backends.megatron import model_provider + from relax.backends.megatron import model_provider, p3o_step if _STATE.args is not None: raise RuntimeError("Task40 forward dump configure() was called more than once in one process") @@ -518,13 +530,20 @@ def configure(args: Any) -> None: _STATE.tp_world_size = mpu.get_tensor_model_parallel_world_size() _STATE.pp_rank = mpu.get_pipeline_model_parallel_rank() _STATE.pp_world_size = mpu.get_pipeline_model_parallel_world_size() + denominator = _STATE.dp_world_size * int(args.micro_batch_size) + if denominator <= 0 or int(args.global_batch_size) % denominator != 0: + raise ValueError( + "Task40 forward dump cannot derive the per-rank capture count from " + f"global_batch_size={args.global_batch_size}, dp_world_size={_STATE.dp_world_size}, " + f"micro_batch_size={args.micro_batch_size}" + ) + _STATE.capture_target_count = int(args.global_batch_size) // denominator _STATE.dump_dir = Path(args.dump_details).parent / "dump" _STATE.dump_dir.mkdir(parents=True, exist_ok=True) (_STATE.dump_dir / f"rank{_STATE.rank:05d}").mkdir(exist_ok=True) _json_write(_STATE.dump_dir / f"runtime_rank{_STATE.rank}.json", _runtime_metadata(args)) - _STATE.original_get_batch = model_backend.get_batch - model_backend.get_batch = _observed_get_batch + _install_get_batch_observers(model_backend, p3o_step) _STATE.original_cp_probe = model_provider._install_cp_probe @@ -551,9 +570,13 @@ def before_train_step( optimizer: Any, opt_param_scheduler: Any, ) -> None: - """Stop capture before the replayed P3O stats/train forwards.""" - del args, rollout_id, step_id, model, optimizer, opt_param_scheduler - _STATE.phase = "train" + """Capture the first P3O forward when rollout log-probs skip actor + forward.""" + del rollout_id, step_id, model, optimizer, opt_param_scheduler + if not _STATE.capture_directories and bool(getattr(args, "use_rollout_logprobs", False)): + _STATE.phase = CAPTURE_PHASE + else: + _STATE.phase = "train" def finalize_manifest(run_dir: Path) -> dict[str, Any]: diff --git a/tests/scripts/local_p3o_task40/test_p3o_forward_dump.py b/tests/scripts/local_p3o_task40/test_p3o_forward_dump.py new file mode 100644 index 000000000..c1b15c85a --- /dev/null +++ b/tests/scripts/local_p3o_task40/test_p3o_forward_dump.py @@ -0,0 +1,57 @@ +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from scripts.local_p3o_task40 import p3o_forward_dump # noqa: E402 + + +def test_install_get_batch_observers_updates_both_imported_aliases(monkeypatch: pytest.MonkeyPatch) -> None: + def original_get_batch() -> None: + return None + + state = p3o_forward_dump.DumpState() + monkeypatch.setattr(p3o_forward_dump, "_STATE", state) + model_backend = SimpleNamespace(get_batch=original_get_batch) + p3o_step_backend = SimpleNamespace(get_batch=original_get_batch) + + p3o_forward_dump._install_get_batch_observers(model_backend, p3o_step_backend) + + assert state.original_get_batch is original_get_batch + assert model_backend.get_batch is p3o_forward_dump._observed_get_batch + assert p3o_step_backend.get_batch is p3o_forward_dump._observed_get_batch + + +def test_install_get_batch_observers_rejects_mismatched_original_aliases() -> None: + model_backend = SimpleNamespace(get_batch=lambda: None) + p3o_step_backend = SimpleNamespace(get_batch=lambda: None) + + with pytest.raises(RuntimeError, match="do not share"): + p3o_forward_dump._install_get_batch_observers(model_backend, p3o_step_backend) + + +@pytest.mark.parametrize( + ("use_rollout_logprobs", "captures", "expected_phase"), + [ + (True, [], p3o_forward_dump.CAPTURE_PHASE), + (True, ["rank00000/micro00000"], "train"), + (False, [], "train"), + ], +) +def test_before_train_step_selects_stats_forward_only_when_actor_forward_is_skipped( + monkeypatch: pytest.MonkeyPatch, + use_rollout_logprobs: bool, + captures: list[str], + expected_phase: str, +) -> None: + state = p3o_forward_dump.DumpState(capture_directories=captures) + monkeypatch.setattr(p3o_forward_dump, "_STATE", state) + args = SimpleNamespace(use_rollout_logprobs=use_rollout_logprobs) + + p3o_forward_dump.before_train_step(args, 0, 0, None, None, None) + + assert state.phase == expected_phase From abfe66dc127ab446dca400d32f48ac2428029518 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:24:15 +0800 Subject: [PATCH 15/30] chore(p3o): record CP forward root cause From e16d325996644f6b205d5b50a189d7e3b35d23da Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:15:11 +0800 Subject: [PATCH 16/30] fix(p3o): equalize CP attention kernel order --- examples/algorithms/p3o/README.md | 12 + examples/algorithms/p3o/README_zh.md | 8 + relax/backends/megatron/data.py | 9 + relax/backends/megatron/model.py | 264 +++++++++++++++++- relax/utils/arguments.py | 11 + ...test_p3o_attention_partition_invariance.py | 231 +++++++++++++++ tests/utils/test_p3o_arguments.py | 24 ++ 7 files changed, 558 insertions(+), 1 deletion(-) create mode 100644 tests/backends/megatron/test_p3o_attention_partition_invariance.py diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md index b196c089c..1da991d64 100644 --- a/examples/algorithms/p3o/README.md +++ b/examples/algorithms/p3o/README.md @@ -86,6 +86,18 @@ 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 +self-attention layer 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 the QKV projection and attention kernel see the +same token order and shape as CP1. It duplicates full-sequence QKV, attention, +and projection 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 validated 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. + 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, diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md index ad26218a2..d8b6693b2 100644 --- a/examples/algorithms/p3o/README_zh.md +++ b/examples/algorithms/p3o/README_zh.md @@ -70,6 +70,14 @@ OPD advantage replacement 或 OPD-only reward 都会形成未经验证的混合 reward/verifier 名称 `P3O_RM_TYPE=mopd` 与 `--use-opd` 训练功能无关，兼容数据集 仍可使用。 +当 `--context-parallel-size > 1` 时，P3O 会自动在每层 THD self-attention 之前重建 +全序列，完成计算后再把输出切回各 CP rank。该严格路径对 `micro-batch` 和 `step` +两种 ESS scope 都生效，使 QKV 投影和 attention kernel 看到与 CP1 相同的 token +顺序和形状。它会在每个 CP rank 上重复完整序列的 QKV、attention 和输出投影计算及 +activation，因此峰值显存和计算量都高于原生 CP；长上下文建议开启整层 activation +recomputation。当前已验证合同是标准 zig-zag THD 且 tensor parallel size 为 1。如果全序列路径发生 +OOM，P3O 会终止并拒绝 CP>1，不会静默回退到数值不等价的原生 CP kernel order。 + 正式默认值为 G=16、global batch 64、micro-batch 1、rollout batch 4、response length 4096 和 30 个 optimizer step（`--num-rollout 30`）。计划配对 seed 为 42、 123 和 2026。smoke 使用 G=4、global batch 16、response length 128 和 1 个 diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index 374c3d951..a8d5f8323 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -441,6 +441,7 @@ def get_batch( tokens = F.pad(tokens, (0, pad), value=pad_token_id) cu_seqlens_list.append(cu_seqlens_list[-1] + pad) + cu_seqlens_cpu = cu_seqlens_list cu_seqlens = torch.tensor( cu_seqlens_list, dtype=torch.int, device=device_utils.make_current_torch_device() ) @@ -464,6 +465,7 @@ def get_batch( cu_seqlens.append(cu_seqlens[-1] + pad) # thd requires the cu_seqlens to be of the origin length + cu_seqlens_cpu = [value * cp_size for value in cu_seqlens] cu_seqlens = ( torch.tensor(cu_seqlens, dtype=torch.int).to(device_utils.make_current_torch_device()) * cp_size ) @@ -479,6 +481,13 @@ def get_batch( if use_dynamic_context_parallel: packed_seq_params.local_cp_size = cp_size packed_seq_params.cp_group = cp_group + if getattr(get_args(), "advantage_estimator", None) == "p3o": + # P3O's strict CP attention reconstructs the CP1 QKV shape without + # changing this existing token layout. Keep the host boundaries from + # construction so every layer avoids a device-to-host synchronization. + packed_seq_params._relax_total_lengths = list(batch["total_lengths"]) + packed_seq_params._relax_attention_pad_multiple = pad_size + packed_seq_params._relax_cu_seqlens_cpu = cu_seqlens_cpu tokens = tokens.unsqueeze(0) else: diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 25bd9db32..f2e8e3f6d 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. import contextlib +import copy import dataclasses import gc import math @@ -10,7 +11,7 @@ from argparse import Namespace from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager -from functools import partial +from functools import partial, wraps from pathlib import Path import torch @@ -55,6 +56,260 @@ logger = get_logger(__name__) +P3O_CP_ATTENTION_OOM_ERROR = ( + "P3O strict mode refuses context-parallel size greater than one after full-sequence attention OOM" +) + + +class _P3OSingletonCPGroup: + """Non-collective CP=1 view used by Megatron's THD RoPE path.""" + + @staticmethod + def size() -> int: + return 1 + + @staticmethod + def rank() -> int: + return 0 + + +_P3O_SINGLETON_CP_GROUP = _P3OSingletonCPGroup() + + +def _resolve_p3o_attention_cp(attention: torch.nn.Module, packed_seq_params: object) -> tuple[int, object, int]: + """Resolve the effective CP group for one packed attention forward.""" + local_cp_size = getattr(packed_seq_params, "local_cp_size", None) + if local_cp_size is not None: + cp_group = getattr(packed_seq_params, "cp_group", None) + if cp_group is None: + raise RuntimeError("P3O full-sequence attention requires packed_seq_params.cp_group for dynamic CP") + cp_size = int(local_cp_size) + else: + cp_group = attention.pg_collection.cp + cp_size = int(cp_group.size()) + cp_rank = int(cp_group.rank()) if cp_size > 1 else 0 + return cp_size, cp_group, cp_rank + + +def _p3o_cp_gather_full( + hidden_states: torch.Tensor, + cu_seqlens: list[int], + cp_size: int, + cp_group: object, +) -> torch.Tensor: + """Gather THD hidden-state shards before the shape-sensitive QKV GEMM.""" + from .cp_utils import gdn_cp_gather_full + + return gdn_cp_gather_full(hidden_states, cu_seqlens, cp_size, cp_group) + + +def _p3o_cp_slice( + full_output: torch.Tensor, + cu_seqlens: list[int], + cp_size: int, + cp_rank: int, +) -> torch.Tensor: + """Return the current rank's THD zig-zag shard from full attention + output.""" + from .cp_utils import gdn_cp_slice + + return gdn_cp_slice(full_output, cu_seqlens, cp_size, cp_rank) + + +def _canonicalize_p3o_attention_input( + full_hidden_states: torch.Tensor, + source_cu_seqlens: list[int], + total_lengths: list[int], + pad_multiple: int, +) -> tuple[torch.Tensor, list[int]]: + """Remove CP-only padding and reproduce the CP1 packed sequence shape.""" + if pad_multiple <= 0: + raise ValueError(f"P3O attention pad multiple must be positive, got {pad_multiple}") + if not total_lengths: + raise ValueError("P3O full-sequence attention requires at least one packed sequence") + if len(source_cu_seqlens) not in {len(total_lengths) + 1, len(total_lengths) + 2}: + raise ValueError( + f"P3O attention metadata mismatch: cu_seqlens={len(source_cu_seqlens)}, total_lengths={len(total_lengths)}" + ) + if ( + source_cu_seqlens[0] != 0 + or any(end < start for start, end in zip(source_cu_seqlens[:-1], source_cu_seqlens[1:], strict=True)) + or source_cu_seqlens[-1] != full_hidden_states.shape[0] + ): + raise ValueError( + "P3O attention cu_seqlens must be monotonic and span the reconstructed tensor: " + f"cu_seqlens={source_cu_seqlens}, tokens={full_hidden_states.shape[0]}" + ) + + pieces: list[torch.Tensor] = [] + canonical_cu_seqlens = [0] + for index, total_length in enumerate(total_lengths): + source_start = source_cu_seqlens[index] + source_end = source_cu_seqlens[index + 1] + if total_length < 0 or source_end - source_start < total_length: + raise ValueError( + "P3O attention source segment is shorter than its real sequence: " + f"index={index}, source_length={source_end - source_start}, total_length={total_length}" + ) + pieces.append(full_hidden_states[source_start : source_start + total_length]) + canonical_cu_seqlens.append(canonical_cu_seqlens[-1] + total_length) + + canonical_padding = (-canonical_cu_seqlens[-1]) % pad_multiple + if canonical_padding: + trailing_start = source_cu_seqlens[len(total_lengths)] + trailing_end = source_cu_seqlens[-1] + available_padding = min(trailing_end - trailing_start, canonical_padding) + if available_padding: + pieces.append(full_hidden_states[trailing_start : trailing_start + available_padding]) + missing_padding = canonical_padding - available_padding + if missing_padding: + pieces.append(full_hidden_states.new_zeros((missing_padding, *full_hidden_states.shape[1:]))) + canonical_cu_seqlens.append(canonical_cu_seqlens[-1] + canonical_padding) + + return torch.cat(pieces, dim=0), canonical_cu_seqlens + + +def _restore_p3o_attention_output_layout( + canonical_output: torch.Tensor, + source_cu_seqlens: list[int], + total_lengths: list[int], +) -> torch.Tensor: + """Map real-token outputs back into the original CP-padded THD layout.""" + if canonical_output.shape[0] < sum(total_lengths): + raise ValueError( + "P3O canonical attention output is shorter than the real-token total: " + f"output={canonical_output.shape[0]}, total={sum(total_lengths)}" + ) + output = canonical_output.new_zeros((source_cu_seqlens[-1], *canonical_output.shape[1:])) + canonical_start = 0 + for index, total_length in enumerate(total_lengths): + source_start = source_cu_seqlens[index] + output[source_start : source_start + total_length] = canonical_output[ + canonical_start : canonical_start + total_length + ] + canonical_start += total_length + return output + + +def _build_p3o_full_sequence_attention_forward(attention: torch.nn.Module) -> Callable: + """Wrap one SelfAttention instance with P3O's strict CP-equivalent path.""" + original_forward = attention.forward + + @wraps(original_forward) + def full_sequence_forward( + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + key_value_states: torch.Tensor | None = None, + inference_context: object | None = None, + rotary_pos_emb: object | None = None, + rotary_pos_cos: torch.Tensor | None = None, + rotary_pos_sin: torch.Tensor | None = None, + rotary_pos_cos_sin: torch.Tensor | None = None, + attention_bias: torch.Tensor | None = None, + packed_seq_params: object | None = None, + sequence_len_offset: int | None = None, + *, + inference_params: object | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + raise RuntimeError("P3O full-sequence attention requires THD packed_seq_params when CP>1") + + cp_size, cp_group, cp_rank = _resolve_p3o_attention_cp(attention, packed_seq_params) + if cp_size == 1: + return original_forward( + hidden_states, + attention_mask, + key_value_states, + inference_context, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + rotary_pos_cos_sin, + attention_bias, + packed_seq_params, + sequence_len_offset, + inference_params=inference_params, + ) + + cu_seqlens_cpu = getattr(packed_seq_params, "_relax_cu_seqlens_cpu", None) + if cu_seqlens_cpu is None: + cu_seqlens_cpu = packed_seq_params.cu_seqlens_q.tolist() + packed_seq_params._relax_cu_seqlens_cpu = cu_seqlens_cpu + + original_cp_group = attention.pg_collection.cp + try: + full_hidden_states = _p3o_cp_gather_full(hidden_states, cu_seqlens_cpu, cp_size, cp_group) + total_lengths = getattr(packed_seq_params, "_relax_total_lengths", None) + pad_multiple = getattr(packed_seq_params, "_relax_attention_pad_multiple", None) + if total_lengths is None or pad_multiple is None: + raise RuntimeError("P3O full-sequence attention requires Relax packed-length metadata") + canonical_hidden_states, canonical_cu_seqlens = _canonicalize_p3o_attention_input( + full_hidden_states, + cu_seqlens_cpu, + total_lengths, + pad_multiple, + ) + full_packed_seq_params = copy.copy(packed_seq_params) + full_packed_seq_params.local_cp_size = 1 + full_packed_seq_params.cp_group = _P3O_SINGLETON_CP_GROUP + canonical_cu_seqlens_tensor = packed_seq_params.cu_seqlens_q.new_tensor(canonical_cu_seqlens) + full_packed_seq_params.cu_seqlens_q = canonical_cu_seqlens_tensor + full_packed_seq_params.cu_seqlens_kv = canonical_cu_seqlens_tensor + full_packed_seq_params.max_seqlen_q = max( + end - start for start, end in zip(canonical_cu_seqlens[:-1], canonical_cu_seqlens[1:], strict=True) + ) + full_packed_seq_params.max_seqlen_kv = full_packed_seq_params.max_seqlen_q + if hasattr(full_packed_seq_params, "cu_seqlens_q_padded"): + full_packed_seq_params.cu_seqlens_q_padded = None + full_packed_seq_params.cu_seqlens_kv_padded = None + full_output, bias = original_forward( + canonical_hidden_states, + attention_mask, + key_value_states, + inference_context, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + rotary_pos_cos_sin, + attention_bias, + full_packed_seq_params, + sequence_len_offset, + inference_params=inference_params, + ) + source_layout_output = _restore_p3o_attention_output_layout( + full_output, + cu_seqlens_cpu, + total_lengths, + ) + return _p3o_cp_slice(source_layout_output, cu_seqlens_cpu, cp_size, cp_rank), bias + except torch.cuda.OutOfMemoryError as exc: + raise RuntimeError(P3O_CP_ATTENTION_OOM_ERROR) from exc + finally: + attention.pg_collection.cp = original_cp_group + + return full_sequence_forward + + +def _install_p3o_full_sequence_attention(args: Namespace, model: torch.nn.Module) -> int: + """Install strict full-sequence attention for P3O whenever configured CP is + greater than one.""" + if getattr(args, "advantage_estimator", None) != "p3o" or getattr(args, "context_parallel_size", 1) <= 1: + return 0 + + installed = 0 + for module in model.modules(): + if type(module).__name__ != "SelfAttention" or getattr( + module, "_relax_p3o_full_sequence_attention_installed", False + ): + continue + if not hasattr(module, "pg_collection") or not hasattr(module.pg_collection, "cp"): + raise RuntimeError("P3O full-sequence attention requires SelfAttention.pg_collection.cp") + module.forward = _build_p3o_full_sequence_attention_forward(module) + module._relax_p3o_full_sequence_attention_installed = True + installed += 1 + return installed + + def _find_lm_output_layer(model: torch.nn.Module) -> torch.nn.Module | None: """Walk DDP / bridge-VL wrappers to the lm_head; None on non-last PP stages. @@ -425,6 +680,13 @@ def setup_model_and_optimizer( wrap_with_ddp=role in ["actor", "critic"], ) + if getattr(args, "advantage_estimator", None) == "p3o" and getattr(args, "context_parallel_size", 1) > 1: + installed_attention_modules = sum( + _install_p3o_full_sequence_attention(args, model_chunk) for model_chunk in model + ) + if installed_attention_modules == 0: + raise RuntimeError("P3O context parallelism requires at least one supported SelfAttention module") + # Some model providers (e.g., Qwen3VLGPTModel) rebuild the decoder in __init__, # which causes duplicate RoutingReplay registrations. Rebuild the list from # the actual model modules to remove stale (orphaned) entries. diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 701f91a54..3bbc76973 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -2969,6 +2969,17 @@ def _validate_p3o_args(args: argparse.Namespace) -> None: if clip_low < 0.0 or clip_high < 0.0: raise ValueError(f"--clip-low/--clip-high must be non-negative, got {clip_low}, {clip_high}.") + context_parallel_size = getattr(args, "context_parallel_size", 1) + if context_parallel_size > 1: + if getattr(args, "qkv_format", "thd") != "thd": + raise ValueError("P3O context parallelism requires THD full-sequence attention.") + if getattr(args, "tensor_model_parallel_size", 1) != 1: + raise ValueError("P3O strict full-sequence attention currently requires --tensor-model-parallel-size 1.") + if getattr(args, "is_vl_model", False): + raise ValueError("P3O strict full-sequence context parallelism currently supports text-only THD models.") + if getattr(args, "allgather_cp", False): + raise ValueError("P3O strict full-sequence attention requires standard zig-zag THD context parallelism.") + if not args.use_rollout_logprobs: raise ValueError( "P3O requires the rollout sampling distribution as its behavior policy. " diff --git a/tests/backends/megatron/test_p3o_attention_partition_invariance.py b/tests/backends/megatron/test_p3o_attention_partition_invariance.py new file mode 100644 index 000000000..bb7bf277a --- /dev/null +++ b/tests/backends/megatron/test_p3o_attention_partition_invariance.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU CP2 parity tests for P3O full-sequence attention.""" + +from __future__ import annotations + +import os +import socket +from argparse import Namespace +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from relax.backends.megatron.cp_utils import gdn_cp_slice +from relax.backends.megatron.model import ( + P3O_CP_ATTENTION_OOM_ERROR, + _install_p3o_full_sequence_attention, +) + + +class _AttentionRoot(torch.nn.Module): + def __init__(self, attention: torch.nn.Module): + super().__init__() + self.attention = attention + + +class SelfAttention(torch.nn.Module): + """Minimal Megatron-shaped attention used to exercise the adapter.""" + + def __init__(self, cp_group: object, weight: torch.Tensor): + super().__init__() + self.pg_collection = SimpleNamespace(cp=cp_group) + self.weight = torch.nn.Parameter(weight.clone()) + self.seen_local_cp_sizes: list[int | None] = [] + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + key_value_states: torch.Tensor | None = None, + inference_context: object | None = None, + rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_cos: torch.Tensor | None = None, + rotary_pos_sin: torch.Tensor | None = None, + rotary_pos_cos_sin: torch.Tensor | None = None, + attention_bias: torch.Tensor | None = None, + packed_seq_params: object | None = None, + sequence_len_offset: int | None = None, + *, + inference_params: object | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + del ( + attention_mask, + key_value_states, + inference_context, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + rotary_pos_cos_sin, + attention_bias, + sequence_len_offset, + inference_params, + ) + self.seen_local_cp_sizes.append(getattr(packed_seq_params, "local_cp_size", None)) + projected = hidden_states @ self.weight + cu_seqlens = packed_seq_params.cu_seqlens_q.tolist() + output = torch.cat( + [ + torch.cumsum(projected[start:end], dim=0) + for start, end in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=True) + ], + dim=0, + ) + return output, torch.zeros(projected.shape[-1], dtype=projected.dtype) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _p3o_args(scope: str, **overrides: object) -> Namespace: + values = { + "advantage_estimator": "p3o", + "context_parallel_size": 2, + "p3o_ess_scope": scope, + } + values.update(overrides) + return Namespace(**values) + + +def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + try: + torch.manual_seed(17) + total_lengths = [5, 5] + cp_padded_lengths = [8, 8] + cu_seqlens = [0, 8, 16] + canonical_cu_seqlens = [0, 5, 10, 16] + hidden_size = 3 + real_samples = [torch.randn(length, 1, hidden_size, dtype=torch.float64) for length in total_lengths] + cp_padded_samples = [ + torch.cat([sample, torch.zeros(padded - len(sample), 1, hidden_size, dtype=torch.float64)]) + for sample, padded in zip(real_samples, cp_padded_lengths, strict=True) + ] + full_hidden = torch.cat(cp_padded_samples) + canonical_hidden = torch.cat([*real_samples, torch.zeros(6, 1, hidden_size, dtype=torch.float64)]) + weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) + + reference_hidden = canonical_hidden.clone().requires_grad_(True) + reference_weight = weight.clone().requires_grad_(True) + reference_projected = reference_hidden @ reference_weight + reference_output = torch.cat( + [ + torch.cumsum(reference_projected[start:end], dim=0) + for start, end in zip(canonical_cu_seqlens[:-1], canonical_cu_seqlens[1:], strict=True) + ] + ) + reference_mask = torch.zeros(len(reference_output), 1, 1, dtype=torch.float64) + reference_mask[: sum(total_lengths)] = 1.0 + (reference_output * reference_mask).square().sum().backward() + + local_hidden = gdn_cp_slice(full_hidden, cu_seqlens, world_size, rank).clone().requires_grad_(True) + attention = SelfAttention(dist.group.WORLD, weight) + root = _AttentionRoot(attention) + assert _install_p3o_full_sequence_attention(_p3o_args(scope), root) == 1 + + packed_seq_params = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor(cu_seqlens, dtype=torch.int32), + cu_seqlens_kv=torch.tensor(cu_seqlens, dtype=torch.int32), + local_cp_size=None, + cp_group=None, + _relax_total_lengths=total_lengths, + _relax_attention_pad_multiple=8, + _relax_cu_seqlens_cpu=cu_seqlens, + ) + output, bias = attention(local_hidden, None, packed_seq_params=packed_seq_params) + expected_full_output = torch.zeros_like(full_hidden) + source_offset = 0 + canonical_offset = 0 + for total_length, padded_length in zip(total_lengths, cp_padded_lengths, strict=True): + expected_full_output[source_offset : source_offset + total_length] = reference_output.detach()[ + canonical_offset : canonical_offset + total_length + ] + source_offset += padded_length + canonical_offset += total_length + expected_output = gdn_cp_slice(expected_full_output, cu_seqlens, world_size, rank) + assert torch.equal(output.detach(), expected_output) + assert torch.equal(bias, torch.zeros(hidden_size, dtype=torch.float64)) + assert attention.seen_local_cp_sizes == [1] + + full_loss_mask = torch.zeros(len(full_hidden), 1, 1, dtype=torch.float64) + source_offset = 0 + for total_length, padded_length in zip(total_lengths, cp_padded_lengths, strict=True): + full_loss_mask[source_offset : source_offset + total_length] = 1.0 + source_offset += padded_length + local_loss_mask = gdn_cp_slice(full_loss_mask, cu_seqlens, world_size, rank) + (output * local_loss_mask).square().sum().backward() + expected_full_input_grad = torch.zeros_like(full_hidden) + source_offset = 0 + canonical_offset = 0 + for total_length, padded_length in zip(total_lengths, cp_padded_lengths, strict=True): + expected_full_input_grad[source_offset : source_offset + total_length] = reference_hidden.grad[ + canonical_offset : canonical_offset + total_length + ] + source_offset += padded_length + canonical_offset += total_length + expected_input_grad = gdn_cp_slice(expected_full_input_grad, cu_seqlens, world_size, rank) + torch.testing.assert_close(local_hidden.grad, expected_input_grad, atol=1e-10, rtol=1e-10) + dist.all_reduce(attention.weight.grad, op=dist.ReduceOp.SUM) + torch.testing.assert_close(attention.weight.grad, reference_weight.grad, atol=1e-10, rtol=1e-10) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("scope", ["micro-batch", "step"]) +def test_p3o_full_sequence_attention_cpu_cp2_matches_cp1(scope: str) -> None: + mp.spawn(_cp2_parity_worker, args=(2, _free_port(), scope), nprocs=2, join=True) + + +def test_p3o_full_sequence_attention_gate_is_p3o_cp_only() -> None: + attention = SelfAttention(SimpleNamespace(size=lambda: 1, rank=lambda: 0), torch.eye(2)) + root = _AttentionRoot(attention) + original_forward = attention.forward + + assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch", context_parallel_size=1), root) == 0 + assert attention.forward == original_forward + assert ( + _install_p3o_full_sequence_attention( + _p3o_args("micro-batch", advantage_estimator="grpo", context_parallel_size=2), root + ) + == 0 + ) + assert attention.forward == original_forward + assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 1 + installed_forward = attention.forward + assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 0 + assert attention.forward == installed_forward + + +def test_p3o_full_sequence_attention_oom_refuses_native_cp_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + from relax.backends.megatron import model as model_module + + cp_group = SimpleNamespace(size=lambda: 2, rank=lambda: 0) + attention = SelfAttention(cp_group, torch.eye(2)) + root = _AttentionRoot(attention) + assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 1 + + def raise_oom(*args: object, **kwargs: object) -> torch.Tensor: + del args, kwargs + raise torch.cuda.OutOfMemoryError("synthetic OOM") + + monkeypatch.setattr(model_module, "_p3o_cp_gather_full", raise_oom) + packed_seq_params = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor([0, 8], dtype=torch.int32), + cu_seqlens_kv=torch.tensor([0, 8], dtype=torch.int32), + local_cp_size=None, + cp_group=None, + ) + with pytest.raises(RuntimeError, match=P3O_CP_ATTENTION_OOM_ERROR): + attention(torch.randn(4, 1, 2), None, packed_seq_params=packed_seq_params) + assert attention.seen_local_cp_sizes == [] + assert attention.pg_collection.cp is cp_group diff --git a/tests/utils/test_p3o_arguments.py b/tests/utils/test_p3o_arguments.py index d94bd250f..7ea3eb3de 100644 --- a/tests/utils/test_p3o_arguments.py +++ b/tests/utils/test_p3o_arguments.py @@ -106,6 +106,11 @@ def _p3o_args(**overrides) -> Namespace: use_routing_replay=False, use_rollout_routing_replay=False, overlap_moe_expert_parallel_comm=False, + context_parallel_size=1, + tensor_model_parallel_size=1, + qkv_format="thd", + is_vl_model=False, + allgather_cp=False, ) config.update(overrides) return Namespace(**config) @@ -182,6 +187,25 @@ def test_p3o_arguments_step_scope_allows_recomputed_loss(): validate_p3o_args(_p3o_args(p3o_ess_scope="step", recompute_loss_function=True)) +@pytest.mark.parametrize("scope", ["micro-batch", "step"]) +def test_p3o_arguments_accepts_strict_full_sequence_cp(scope): + validate_p3o_args(_p3o_args(p3o_ess_scope=scope, context_parallel_size=2)) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"qkv_format": "bshd"}, "THD"), + ({"tensor_model_parallel_size": 2}, "tensor-model-parallel-size 1"), + ({"is_vl_model": True}, "text-only THD"), + ({"allgather_cp": True}, "standard zig-zag THD"), + ], +) +def test_p3o_arguments_rejects_unsupported_strict_cp_attention(overrides, message): + with pytest.raises(ValueError, match=message): + validate_p3o_args(_p3o_args(context_parallel_size=2, **overrides)) + + def test_megatron_checkpoint_source_recognizes_root_and_direct_iteration_paths(tmp_path): checkpoint_root = tmp_path / "checkpoint" checkpoint_root.mkdir() From a930d73ecd7a958f2306ab0b1483bb756f282992 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:53:39 +0800 Subject: [PATCH 17/30] chore(p3o): record strict step-0 revalidation --- .../analyze_step0_revalidation.py | 542 ++++++++++++++++++ .../run_step0_revalidation.sh | 109 ++++ .../local_p3o_task40/step0_revalidation.py | 424 ++++++++++++++ .../test_step0_revalidation.py | 108 ++++ 4 files changed, 1183 insertions(+) create mode 100755 scripts/local_p3o_task40/analyze_step0_revalidation.py create mode 100755 scripts/local_p3o_task40/run_step0_revalidation.sh create mode 100755 scripts/local_p3o_task40/step0_revalidation.py create mode 100644 tests/scripts/local_p3o_task40/test_step0_revalidation.py diff --git a/scripts/local_p3o_task40/analyze_step0_revalidation.py b/scripts/local_p3o_task40/analyze_step0_revalidation.py new file mode 100755 index 000000000..01ab821ea --- /dev/null +++ b/scripts/local_p3o_task40/analyze_step0_revalidation.py @@ -0,0 +1,542 @@ +#!/usr/bin/env python3 + +"""Apply the frozen Batch-7 three-stage Step-0 acceptance gates.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import torch + +from scripts.local_p3o_task40.analyze_step0_oracle import _load_run, _token_key +from scripts.local_p3o_task40.step0_revalidation import EXPECTED_FIXTURE_SHA256 + + +REL_TOL = 1e-6 +NEAR_ZERO_ABS_TOL = 1e-9 +LOGPROB_ABS_TOL = 1e-6 +GRAD_REL_L2_TOL = 1e-6 +GRAD_COSINE_MIN = 1.0 - 1e-9 +TOPOLOGIES = ("dp1", "dp4cp1", "dp2cp1", "dp2cp2") +METRIC_KEYS = { + "normalized_ess": "train/p3o/normalized_ess", + "adaptive_cap": "train/p3o/adaptive_cap", + "ratio_mean": "train/p3o/ratio_mean", + "ratio_std": "train/p3o/ratio_std", + "cap_fraction": "train/p3o/cap_fraction", + "clip_fraction": "train/p3o/clip_fraction", + "behavior_kl_proxy": "train/p3o/behavior_kl_proxy", + "adaptive_kl_loss": "train/p3o/adaptive_kl_loss", + "reference_kl": "train/p3o/reference_kl", + "score_loss": "train/p3o/score_loss", + "entropy": "train/p3o/entropy", + "total_loss": "train/p3o/total_loss", + "loss": "train/loss", +} + + +def _json_write(path: Path, payload: Any) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + temporary.replace(path) + + +def _plain(value: Any) -> Any: + if torch.is_tensor(value): + return value.detach().cpu().tolist() + if isinstance(value, (list, tuple)): + return [_plain(item) for item in value] + if isinstance(value, dict): + return {str(key): _plain(item) for key, item in value.items()} + return value + + +def _canonical_cu(total_lengths: list[int], pad_multiple: int) -> list[int]: + boundaries = [0] + for length in total_lengths: + boundaries.append(boundaries[-1] + int(length)) + padding = (-boundaries[-1]) % pad_multiple + if padding: + boundaries.append(boundaries[-1] + padding) + return boundaries + + +def _expected_source_cu(total_lengths: list[int], cp_size: int, pad_multiple: int) -> list[int]: + boundaries = [0] + for length in total_lengths: + padded = int(length) if cp_size == 1 else math.ceil(int(length) / (2 * cp_size)) * (2 * cp_size) + boundaries.append(boundaries[-1] + padded) + local_length = boundaries[-1] // cp_size + local_padding = (-local_length) % pad_multiple + if local_padding: + boundaries.append(boundaries[-1] + local_padding * cp_size) + return boundaries + + +def _sample_metadata(run_dir: Path) -> tuple[dict[str, Any], dict[str, Any]]: + runtimes = { + int(record["rank"]): record + for record in ( + json.loads(path.read_text()) for path in sorted((run_dir / "oracle").glob("runtime_rank*.json")) + ) + } + samples: dict[str, Any] = {} + cu_records: dict[str, Any] = {} + for path in sorted((run_dir / "oracle").glob("vectors_rank*_micro*.pt")): + artifact = torch.load(path, map_location="cpu", weights_only=False) + runtime = runtimes[int(artifact["rank"])] + cp_size = int(runtime["cp_world_size"]) + tokens_list = artifact["token_ids"] + positions_list = artifact["position_ids"] + loss_masks = artifact["loss_masks"] + total_lengths = [int(value) for value in artifact["total_lengths"]] + response_lengths = [int(value) for value in artifact["response_lengths"]] + if len(tokens_list) != len(total_lengths): + raise ValueError(f"token/length count mismatch in {path}") + for index, (tokens, positions, total_length, response_length) in enumerate( + zip(tokens_list, positions_list, total_lengths, response_lengths, strict=True) + ): + key = _token_key(tokens) + loss_mask = loss_masks[index] if isinstance(loss_masks, (list, tuple)) else loss_masks + record = { + "token_ids": _plain(tokens.to(torch.int64)), + "position_ids": _plain(positions.to(torch.int64)), + "total_length": total_length, + "response_length": response_length, + "loss_mask": _plain(loss_mask), + } + if key in samples and samples[key] != record: + raise ValueError(f"inconsistent per-token metadata for sample {key} in {run_dir}") + samples[key] = record + + direct_q = _plain(artifact.get("cu_seqlens_q")) + direct_kv = _plain(artifact.get("cu_seqlens_kv")) + pad_multiple = int(artifact.get("relax_attention_pad_multiple") or 0) + if direct_q is None or direct_kv is None or pad_multiple <= 0: + raise ValueError(f"missing cu_seqlens or P3O pad metadata in {path}") + sample_group_key = "+".join(_token_key(tokens) for tokens in tokens_list) + canonical = _canonical_cu(total_lengths, pad_multiple) + expected_source = _expected_source_cu(total_lengths, cp_size, pad_multiple) + relation = { + "sample_keys": sample_group_key.split("+"), + "cp_size": cp_size, + "pad_multiple": pad_multiple, + "total_lengths": total_lengths, + "direct_cu_seqlens_q": direct_q, + "direct_cu_seqlens_kv": direct_kv, + "relax_cu_seqlens_cpu": _plain(artifact.get("relax_cu_seqlens_cpu")), + "expected_source_cu_seqlens": expected_source, + "derived_canonical_cp1_cu_seqlens": canonical, + "source_relation_pass": direct_q == direct_kv == expected_source, + "derivation": ( + "CP1 keeps each real sample length unchanged. CP>1 pads each sample to 2*CP before zig-zag " + "slicing; rank-local concatenation is then padded to pad_multiple and source boundaries are " + "multiplied by CP. The Batch-6 adapter removes CP-only padding, concatenates total_lengths, " + "and appends only CP1 tail padding." + ), + } + if sample_group_key in cu_records and cu_records[sample_group_key] != relation: + raise ValueError(f"inconsistent cu_seqlens relation for {sample_group_key} in {run_dir}") + cu_records[sample_group_key] = relation + return samples, cu_records + + +def _numeric_comparison(reference: float, candidate: float) -> dict[str, Any]: + absolute_error = abs(reference - candidate) + scale = max(abs(reference), abs(candidate)) + relative_error = absolute_error / max(scale, 1e-30) + near_zero = scale <= NEAR_ZERO_ABS_TOL + passed = absolute_error <= NEAR_ZERO_ABS_TOL if near_zero else relative_error <= REL_TOL + return { + "reference": reference, + "candidate": candidate, + "absolute_error": absolute_error, + "relative_error": relative_error, + "near_zero_rule": near_zero, + "threshold": NEAR_ZERO_ABS_TOL if near_zero else REL_TOL, + "pass": passed, + } + + +def _load_parameter_manifests(run_dir: Path) -> list[dict[str, Any]]: + return [ + json.loads(path.read_text()) for path in sorted((run_dir / "oracle").glob("initial_parameters_rank*.json")) + ] + + +def _load_gradient_summaries(run_dir: Path) -> list[dict[str, Any]]: + return [ + json.loads(path.read_text()) for path in sorted((run_dir / "oracle" / "gradients").glob("summary_rank*.json")) + ] + + +def _gradient_shard_map(run_dir: Path, summaries: list[dict[str, Any]]) -> dict[str, list[dict[str, Any]]]: + gradient_root = run_dir / "oracle" / "gradients" + shards: dict[str, list[dict[str, Any]]] = {} + for summary in summaries: + for record in summary["tensors"]: + if not record.get("present"): + continue + relative = record.get("file") + if not relative: + raise ValueError(f"gradient shard is missing a file for {record['name']} in {run_dir}") + path = gradient_root / relative + if not path.is_file(): + raise ValueError(f"missing gradient tensor file: {path}") + enriched = dict(record) + enriched["path"] = path + enriched["rank"] = int(summary["rank"]) + shards.setdefault(record["name"], []).append(enriched) + return shards + + +def _gradient_coverage(run_dir: Path, summaries: list[dict[str, Any]]) -> dict[str, Any]: + parameter_manifests = _load_parameter_manifests(run_dir) + expected = { + record["name"]: int(record["numel"]) + for record in parameter_manifests[0]["tensors"] + if bool(record["requires_grad"]) + } + shards = _gradient_shard_map(run_dir, summaries) + errors = [] + if set(shards) != set(expected): + errors.append( + { + "missing_parameters": sorted(set(expected) - set(shards)), + "unexpected_parameters": sorted(set(shards) - set(expected)), + } + ) + for name in sorted(set(expected) & set(shards)): + cursor = 0 + for record in sorted(shards[name], key=lambda item: (int(item["range_start"]), int(item["range_end"]))): + start = int(record["range_start"]) + end = int(record["range_end"]) + if start != cursor or end <= start or int(record["numel"]) != end - start: + errors.append( + {"name": name, "expected_start": cursor, "record": {**record, "path": str(record["path"])}} + ) + break + cursor = end + if cursor != expected[name]: + errors.append({"name": name, "covered_until": cursor, "parameter_numel": expected[name]}) + return { + "pass": not errors, + "errors": errors, + "parameter_count": len(expected), + "covered_parameter_count": len(shards), + "aggregate_parameter_numel": sum(expected.values()), + "aggregate_shard_numel": sum(int(record["numel"]) for records in shards.values() for record in records), + } + + +def _reconstruct_gradient(records: list[dict[str, Any]]) -> torch.Tensor: + ordered = sorted(records, key=lambda item: (int(item["range_start"]), int(item["range_end"]))) + parameter_numel = int(ordered[0]["parameter_numel"]) + gradient = torch.empty(parameter_numel, dtype=torch.float32) + cursor = 0 + for record in ordered: + start = int(record["range_start"]) + end = int(record["range_end"]) + if start != cursor or end <= start: + raise ValueError( + f"non-contiguous gradient shards for {record['name']}: expected {cursor}, got {start}:{end}" + ) + shard = torch.load(record["path"], map_location="cpu", weights_only=False).reshape(-1).to(torch.float32) + if shard.numel() != end - start: + raise ValueError(f"gradient shard/file length mismatch for {record['name']}") + gradient[start:end] = shard + cursor = end + if cursor != parameter_numel: + raise ValueError(f"incomplete gradient reconstruction for {ordered[0]['name']}: {cursor}/{parameter_numel}") + return gradient + + +def _compare_gradients(reference_dir: Path, candidate_dir: Path) -> dict[str, Any]: + reference_summaries = _load_gradient_summaries(reference_dir) + candidate_summaries = _load_gradient_summaries(candidate_dir) + reference_shards = _gradient_shard_map(reference_dir, reference_summaries) + candidate_shards = _gradient_shard_map(candidate_dir, candidate_summaries) + names_equal = set(reference_shards) == set(candidate_shards) + if not names_equal: + return { + "pass": False, + "parameter_names_equal": False, + "missing_from_candidate": sorted(set(reference_shards) - set(candidate_shards)), + "unexpected_in_candidate": sorted(set(candidate_shards) - set(reference_shards)), + } + + diff_sq = 0.0 + reference_sq = 0.0 + candidate_sq = 0.0 + dot = 0.0 + per_tensor = [] + for name in sorted(reference_shards): + reference = _reconstruct_gradient(reference_shards[name]) + candidate = _reconstruct_gradient(candidate_shards[name]) + if reference.shape != candidate.shape: + per_tensor.append({"name": name, "shape_equal": False, "pass": False}) + continue + tensor_diff_sq = 0.0 + tensor_reference_sq = 0.0 + tensor_candidate_sq = 0.0 + tensor_dot = 0.0 + for reference_chunk, candidate_chunk in zip( + reference.split(8 * 1024 * 1024), candidate.split(8 * 1024 * 1024), strict=True + ): + reference64 = reference_chunk.to(torch.float64) + candidate64 = candidate_chunk.to(torch.float64) + delta = candidate64 - reference64 + tensor_diff_sq += float(torch.sum(delta * delta)) + tensor_reference_sq += float(torch.sum(reference64 * reference64)) + tensor_candidate_sq += float(torch.sum(candidate64 * candidate64)) + tensor_dot += float(torch.sum(reference64 * candidate64)) + tensor_rel_l2 = math.sqrt(tensor_diff_sq) / max(math.sqrt(tensor_reference_sq), 1e-30) + if tensor_reference_sq == 0.0 and tensor_candidate_sq == 0.0: + tensor_cosine = 1.0 + else: + tensor_cosine = tensor_dot / max(math.sqrt(tensor_reference_sq * tensor_candidate_sq), 1e-30) + tensor_pass = tensor_rel_l2 <= GRAD_REL_L2_TOL and tensor_cosine >= GRAD_COSINE_MIN + per_tensor.append( + { + "name": name, + "shape_equal": True, + "relative_l2": tensor_rel_l2, + "cosine": tensor_cosine, + "pass": tensor_pass, + } + ) + diff_sq += tensor_diff_sq + reference_sq += tensor_reference_sq + candidate_sq += tensor_candidate_sq + dot += tensor_dot + + relative_l2 = math.sqrt(diff_sq) / max(math.sqrt(reference_sq), 1e-30) + cosine = 1.0 if reference_sq == candidate_sq == 0.0 else dot / max(math.sqrt(reference_sq * candidate_sq), 1e-30) + all_shapes = all(record.get("shape_equal", False) for record in per_tensor) + return { + "parameter_names_equal": names_equal, + "all_shapes_equal": all_shapes, + "full_parameter_relative_l2": relative_l2, + "full_parameter_cosine": cosine, + "relative_l2_threshold": GRAD_REL_L2_TOL, + "cosine_minimum": GRAD_COSINE_MIN, + "per_tensor": per_tensor, + "worst_per_tensor_relative_l2": max( + (record.get("relative_l2", math.inf) for record in per_tensor), default=0.0 + ), + "minimum_per_tensor_cosine": min((record.get("cosine", -math.inf) for record in per_tensor), default=1.0), + "pass": all_shapes and relative_l2 <= GRAD_REL_L2_TOL and cosine >= GRAD_COSINE_MIN, + } + + +def _artifact_contract(name: str, run_dir: Path, run: dict[str, Any]) -> dict[str, Any]: + runtimes = run["runtime"] + parameters = _load_parameter_manifests(run_dir) + gradients = _load_gradient_summaries(run_dir) + gradient_coverage = _gradient_coverage(run_dir, gradients) if gradients else {"pass": False, "errors": ["none"]} + expected_world_size = int(runtimes[0]["world_size"]) if runtimes else 0 + expected_vectors = sum( + int(runtime["global_batch_size"]) // (int(runtime["dp_world_size"]) * int(runtime["micro_batch_size"])) + for runtime in runtimes + ) + gradient_hashes = sorted({record["gradient_sha256"] for record in gradients}) + parameter_hashes = sorted({record["parameter_sha256"] for record in parameters}) + resolved = { + "topology": name, + "world_size": expected_world_size, + "dp_world_size": sorted({int(runtime["dp_world_size"]) for runtime in runtimes}), + "cp_world_size": sorted({int(runtime["cp_world_size"]) for runtime in runtimes}), + "tp_world_size": sorted({int(runtime["tp_world_size"]) for runtime in runtimes}), + "pp_world_size": sorted({int(runtime["pp_world_size"]) for runtime in runtimes}), + "bf16": sorted({bool(runtime["bf16"]) for runtime in runtimes}), + "fp16": sorted({bool(runtime["fp16"]) for runtime in runtimes}), + "qkv_format": sorted({str(runtime["qkv_format"]) for runtime in runtimes}), + "micro_batch_size": sorted({int(runtime["micro_batch_size"]) for runtime in runtimes}), + "p3o_ess_scope": sorted({str(runtime["p3o_ess_scope"]) for runtime in runtimes}), + "fixture_sha256": sorted({str(runtime["fixture_sha256"]) for runtime in runtimes}), + } + checks = { + "exit_zero": run["exit_code"] == 0, + "runtime_rank_count": len(runtimes) == expected_world_size, + "oracle_vector_count": run["vector_artifact_count"] == expected_vectors, + "global_stat_rank_count": len(list((run_dir / "oracle").glob("global_stats_rank*_sync0.pt"))) + == expected_world_size, + "parameter_manifest_rank_count": len(parameters) == expected_world_size, + "gradient_summary_rank_count": len(gradients) == expected_world_size, + "fixture_sha": resolved["fixture_sha256"] == [EXPECTED_FIXTURE_SHA256], + "bf16_only": resolved["bf16"] == [True] and resolved["fp16"] == [False], + "thd": resolved["qkv_format"] == ["thd"], + "mbs1": resolved["micro_batch_size"] == [1], + "step_scope": resolved["p3o_ess_scope"] == ["step"], + "position_ids_valid": bool(run["position_ids_valid"]), + "global_stats_rank_agreement": bool(run["global_rank_agreement"]), + "parameters_identical_within_topology": len(parameter_hashes) == 1, + "gradient_shard_coverage": bool(gradient_coverage["pass"]), + "no_missing_owned_gradients": all(not record["missing_owned_gradients"] for record in gradients), + "all_gradients_finite": all( + all((not tensor.get("present")) or tensor.get("finite", False) for tensor in record["tensors"]) + for record in gradients + ), + } + return { + "resolved_contract": resolved, + "checks": checks, + "parameter_sha256_values": parameter_hashes, + "gradient_sha256_values": gradient_hashes, + "gradient_shard_coverage": gradient_coverage, + "pass": all(checks.values()), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + for topology in TOPOLOGIES: + parser.add_argument(f"--{topology}", type=Path, required=True) + parser.add_argument("--output-root", type=Path, required=True) + args = parser.parse_args() + run_dirs = {topology: getattr(args, topology).resolve() for topology in TOPOLOGIES} + runs = {topology: _load_run(run_dir) for topology, run_dir in run_dirs.items()} + metadata = {} + cu_relations = {} + contracts = {} + for topology in TOPOLOGIES: + metadata[topology], cu_relations[topology] = _sample_metadata(run_dirs[topology]) + contracts[topology] = _artifact_contract(topology, run_dirs[topology], runs[topology]) + + reference = runs["dp1"] + reference_metadata = metadata["dp1"] + reference_canonical_cu = { + key: record["derived_canonical_cp1_cu_seqlens"] for key, record in cu_relations["dp1"].items() + } + parameter_hashes = {topology: contracts[topology]["parameter_sha256_values"] for topology in TOPOLOGIES} + cross_topology_parameter_sha_pass = ( + all(len(values) == 1 for values in parameter_hashes.values()) + and len({values[0] for values in parameter_hashes.values()}) == 1 + ) + + overall_pass = True + topology_verdicts = {} + for topology in TOPOLOGIES: + run = runs[topology] + key_match = set(reference["canonical"]) == set(run["canonical"]) + current_max_abs = float(torch.max(torch.abs(reference["current"] - run["current"]))) if key_match else math.inf + behavior_max_abs = ( + float(torch.max(torch.abs(reference["behavior"] - run["behavior"]))) if key_match else math.inf + ) + canonical_cu = { + key: record["derived_canonical_cp1_cu_seqlens"] for key, record in cu_relations[topology].items() + } + stage1_checks = { + "artifact_contract": contracts[topology]["pass"], + "valid_token_keys_identical": key_match, + "current_log_probs_max_abs": current_max_abs <= LOGPROB_ABS_TOL, + "rollout_log_probs_exact": behavior_max_abs == 0.0, + "per_token_metadata_equal": metadata[topology] == reference_metadata, + "initial_parameter_sha_cross_topology": cross_topology_parameter_sha_pass, + "source_cu_derivation_valid": all( + record["source_relation_pass"] for record in cu_relations[topology].values() + ), + "derived_canonical_cu_equal": canonical_cu == reference_canonical_cu, + } + stage1 = { + "status": "PASS" if all(stage1_checks.values()) else "FAIL", + "checks": stage1_checks, + "current_log_probs_max_abs": current_max_abs, + "current_log_probs_max_abs_threshold": LOGPROB_ABS_TOL, + "rollout_log_probs_max_abs": behavior_max_abs, + "initial_parameter_sha256": parameter_hashes[topology], + "cross_topology_initial_parameter_sha256": parameter_hashes, + "cu_seqlens": cu_relations[topology], + } + + stage2_values = { + "s1": (float(reference["global_stats"]["s1"]), float(run["global_stats"]["s1"])), + "s2": (float(reference["global_stats"]["s2"]), float(run["global_stats"]["s2"])), + "n": (float(reference["global_stats"]["n"]), float(run["global_stats"]["n"])), + **{ + name: (float(reference["metrics"][metric]), float(run["metrics"][metric])) + for name, metric in METRIC_KEYS.items() + }, + } + stage2_comparisons = { + name: _numeric_comparison(reference_value, candidate_value) + for name, (reference_value, candidate_value) in stage2_values.items() + } + stage2 = { + "status": "PASS" if all(record["pass"] for record in stage2_comparisons.values()) else "FAIL", + "relative_tolerance": REL_TOL, + "near_zero_absolute_tolerance": NEAR_ZERO_ABS_TOL, + "comparisons": stage2_comparisons, + } + + gradient_comparison = _compare_gradients(run_dirs["dp1"], run_dirs[topology]) + stage3 = {"status": "PASS" if gradient_comparison["pass"] else "FAIL", **gradient_comparison} + stages_pass = stage1["status"] == stage2["status"] == stage3["status"] == "PASS" + overall_pass &= stages_pass + topology_verdicts[topology] = { + "topology": topology, + "reference": "dp1", + "run_dir": str(run_dirs[topology]), + "status": "PASS" if stages_pass else "FAIL", + "artifact_contract": contracts[topology], + "stage1": stage1, + "stage2": stage2, + "stage3": stage3, + } + topology_root = args.output_root / topology + topology_root.mkdir(parents=True, exist_ok=True) + _json_write(topology_root / "stage1_verdict.json", stage1) + _json_write(topology_root / "stage2_verdict.json", stage2) + _json_write(topology_root / "stage3_verdict.json", stage3) + _json_write(topology_root / "cell_verdict.json", topology_verdicts[topology]) + + failed_stages = { + topology: [ + stage for stage in ("stage1", "stage2", "stage3") if topology_verdicts[topology][stage]["status"] != "PASS" + ] + for topology in TOPOLOGIES + if topology_verdicts[topology]["status"] != "PASS" + } + input_identity_checks = { + "artifact_contract", + "valid_token_keys_identical", + "rollout_log_probs_exact", + "per_token_metadata_equal", + "initial_parameter_sha_cross_topology", + "source_cu_derivation_valid", + "derived_canonical_cu_equal", + } + input_identity_failed = any( + not verdict["stage1"]["checks"][check] + for verdict in topology_verdicts.values() + for check in input_identity_checks + ) + route = "Batch 5" if input_identity_failed else ("Batch 6" if failed_stages else None) + result = { + "batch": 7, + "commit": "e16d325996644f6b205d5b50a189d7e3b35d23da", + "status": "COMPLETE" if overall_pass else "FAIL_ROUTING", + "allow_batch8": overall_pass, + "thresholds": { + "current_log_probs_max_abs": LOGPROB_ABS_TOL, + "stage2_relative": REL_TOL, + "stage2_near_zero_absolute": NEAR_ZERO_ABS_TOL, + "gradient_relative_l2": GRAD_REL_L2_TOL, + "gradient_cosine_minimum": GRAD_COSINE_MIN, + }, + "failed_stages": failed_stages, + "route": route, + "topologies": topology_verdicts, + } + args.output_root.mkdir(parents=True, exist_ok=True) + _json_write(args.output_root / "BATCH7_VERDICT.json", result) + print(json.dumps({"status": result["status"], "allow_batch8": overall_pass, "failed_stages": failed_stages})) + raise SystemExit(0 if overall_pass else 1) + + +if __name__ == "__main__": + main() diff --git a/scripts/local_p3o_task40/run_step0_revalidation.sh b/scripts/local_p3o_task40/run_step0_revalidation.sh new file mode 100755 index 000000000..af5698518 --- /dev/null +++ b/scripts/local_p3o_task40/run_step0_revalidation.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash + +# Task40 Batch 7 command package: exactly four BF16 THD MBS1 step-scope cells. + +set -euo pipefail + +if [[ "${1:-}" == "--in-container" ]]; then + TOPOLOGY="${2:?usage: run_step0_revalidation.sh --in-container }" + case "${TOPOLOGY}" in + dp1) ACTOR_WORLD_SIZE=1; CONTEXT_PARALLEL_SIZE=1 ;; + dp4cp1) ACTOR_WORLD_SIZE=4; CONTEXT_PARALLEL_SIZE=1 ;; + dp2cp1) ACTOR_WORLD_SIZE=2; CONTEXT_PARALLEL_SIZE=1 ;; + dp2cp2) ACTOR_WORLD_SIZE=4; CONTEXT_PARALLEL_SIZE=2 ;; + *) echo "unsupported topology: ${TOPOLOGY}" >&2; exit 2 ;; + esac + + : "${P3O_STEP0_FIXTURE:?P3O_STEP0_FIXTURE must be set}" + export P3O_ALGORITHM=p3o + export P3O_ENABLE_TEMPERATURE_OVERRIDE=0 + export P3O_UPDATE_WEIGHTS_INTERVAL=1 + + SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" + REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." >/dev/null 2>&1 && pwd)" + source "${REPO_ROOT}/examples/algorithms/p3o/common_a100x4.sh" + eval "$(declare -f P3O_build_args | sed '1s/P3O_build_args/P3O_build_args_base/')" + + replace_train_arg() { + local flag="$1" + local value="$2" + local index + for ((index = 0; index < ${#P3O_TRAIN_ARGS[@]}; index++)); do + if [[ "${P3O_TRAIN_ARGS[index]}" == "${flag}" ]]; then + P3O_TRAIN_ARGS[index + 1]="${value}" + return 0 + fi + done + echo "required argument not found: ${flag}" >&2 + return 3 + } + + P3O_build_args() { + P3O_build_args_base + replace_train_arg --resource "{\"actor\":[1,${ACTOR_WORLD_SIZE}],\"rollout\":[1,4]}" + replace_train_arg --context-parallel-size "${CONTEXT_PARALLEL_SIZE}" + replace_train_arg --lr 1e-6 + P3O_TRAIN_ARGS+=( + --load-debug-rollout-data "${P3O_STEP0_FIXTURE}" + --custom-megatron-init-path scripts.local_p3o_task40.step0_revalidation.configure + --custom-megatron-before-train-step-hook-path scripts.local_p3o_task40.step0_revalidation.before_train_step + --dump-details "${P3O_OUTPUT_ROOT}/${P3O_CONFIG_NAME}/seed_${P3O_SEED}/${P3O_RUN_ID}/debug" + ) + } + + P3O_CONFIG_NAME="${TOPOLOGY}" + P3O_run + exit $? +fi + +TOPOLOGY="${1:?usage: run_step0_revalidation.sh }" +RUN_ID="${2:?usage: run_step0_revalidation.sh }" +case "${TOPOLOGY}" in + dp1|dp4cp1|dp2cp1|dp2cp2) ;; + *) echo "unsupported topology: ${TOPOLOGY}" >&2; exit 2 ;; +esac + +INFRA=/lustre/home/sztu_camdt_zhanghua/jimaomo/infra +REPO="${INFRA}/Relax" +IMAGE="${INFRA}/images/relaxrl-dev-20260715-8325919e.sif" +CAMPAIGN="${INFRA}/Output/task40/task40_cp_forward_diag_20260817_6c7a3d2" +FIXTURE="${INFRA}/Output/task40/task40_p0_cluster_20260814_6c7a3d2/fixtures/step0_rollout.pt" +CELL_DIR="${CAMPAIGN}/runs/step0_revalidation/${TOPOLOGY}" +mkdir -p "${CELL_DIR}/launcher_logs" + +set +e +apptainer exec --nv --bind /lustre:/lustre "${IMAGE}" env \ + P3O_MODE=smoke \ + P3O_MODEL_CONFIG="${REPO}/scripts/local_p3o_task40/qwen2p5_1p5b.sh" \ + P3O_MODEL_ROTARY_BASE=1000000 \ + P3O_MODEL_DIR="${INFRA}/Qwen2.5-1.5B-Instruct" \ + P3O_TRAIN_DATA="${INFRA}/gsm8k/main/train_clean.parquet" \ + P3O_INPUT_KEY=question \ + P3O_LABEL_KEY=answer \ + P3O_RM_TYPE=openr1mm \ + P3O_OUTPUT_ROOT="${CAMPAIGN}/runs/step0_revalidation" \ + P3O_MEGATRON_DIR=/root/Megatron-LM \ + P3O_RAY_DASHBOARD=http://127.0.0.1:8265 \ + P3O_NUM_ROLLOUT=1 \ + P3O_ROLLOUT_BATCH_SIZE=4 \ + P3O_N_SAMPLES=16 \ + P3O_GLOBAL_BATCH_SIZE=64 \ + P3O_MICRO_BATCH_SIZE=1 \ + P3O_MAX_RESPONSE_LEN=4096 \ + P3O_ESS_SCOPE=step \ + P3O_KL_MODE=proxy_safe \ + P3O_SEED=42 \ + P3O_PIPELINE_MODEL_PARALLEL_SIZE=1 \ + P3O_ACTIVATION_RECOMPUTE=0 \ + P3O_LOG_PROBS_CHUNK_SIZE=1024 \ + P3O_ROLLOUT_SHUFFLE=0 \ + P3O_DETERMINISTIC_INFERENCE=1 \ + P3O_CLEAR_RUNTIME_PROXIES=1 \ + P3O_STEP0_FIXTURE="${FIXTURE}" \ + P3O_RUN_ID="${RUN_ID}" \ + bash -lc "cd '${REPO}' && bash scripts/local_p3o_task40/run_step0_revalidation.sh --in-container '${TOPOLOGY}'" \ + 2>&1 | tee "${CELL_DIR}/launcher_logs/${RUN_ID}.log" +EXIT_CODE=${PIPESTATUS[0]} +set -e +printf '%s\n' "${EXIT_CODE}" >"${CELL_DIR}/launcher_logs/${RUN_ID}.exit_code.txt" +exit "${EXIT_CODE}" diff --git a/scripts/local_p3o_task40/step0_revalidation.py b/scripts/local_p3o_task40/step0_revalidation.py new file mode 100755 index 000000000..5af1122c2 --- /dev/null +++ b/scripts/local_p3o_task40/step0_revalidation.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python3 + +"""Batch-7 BF16 Step-0 observability hook. + +The hook is intentionally local to Task40. It records the fixed-input P3O +oracle, the checkpoint-loaded parameter identity, and finalized gradients +without changing the production training path. +""" + +from __future__ import annotations + +import hashlib +import json +import threading +from pathlib import Path +from typing import Any, Iterator + + +EXPECTED_FIXTURE_SHA256 = "48538d165386dc94006613d857c022a7ba2e979bdc31bc617374eee2dc3c35b8" +FORMAT_VERSION = 1 +_STATE: dict[str, Any] = { + "args": None, + "output_dir": None, + "rank": 0, + "local_counter": 0, + "sync_counter": 0, + "parameters_captured": False, + "parameter_capture_error": None, + "parameter_capture_thread": None, + "gradients_captured": False, + "optimizer_ids": set(), +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _write_json(path: Path, payload: Any) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + temporary.replace(path) + + +def _cpu(value: Any) -> Any: + import torch + + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, list): + return [_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(_cpu(item) for item in value) + return value + + +def _tensor_bytes(tensor: Any) -> tuple[Any, bytes]: + import torch + + cpu = tensor.detach().cpu().contiguous() + raw = cpu.view(torch.uint8).numpy().tobytes() + return cpu, raw + + +def _named_parameters(model: Any) -> Iterator[tuple[str, Any]]: + seen: set[int] = set() + for chunk_index, chunk in enumerate(model): + for name, parameter in chunk.named_parameters(): + if id(parameter) in seen: + continue + seen.add(id(parameter)) + yield f"chunk{chunk_index:03d}.{name}", parameter + + +def _capture_initial_parameters(model: Any) -> None: + if _STATE["parameters_captured"]: + return + output_dir: Path = _STATE["output_dir"] + rank = int(_STATE["rank"]) + aggregate = hashlib.sha256() + tensors = [] + for name, parameter in _named_parameters(model): + cpu, raw = _tensor_bytes(parameter) + header = json.dumps( + {"dtype": str(cpu.dtype), "name": name, "shape": list(cpu.shape)}, + sort_keys=True, + separators=(",", ":"), + ).encode() + aggregate.update(len(header).to_bytes(8, "little")) + aggregate.update(header) + aggregate.update(len(raw).to_bytes(8, "little")) + aggregate.update(raw) + tensors.append( + { + "name": name, + "shape": list(cpu.shape), + "dtype": str(cpu.dtype), + "numel": cpu.numel(), + "sha256": hashlib.sha256(raw).hexdigest(), + "requires_grad": bool(parameter.requires_grad), + } + ) + _write_json( + output_dir / f"initial_parameters_rank{rank}.json", + { + "format_version": FORMAT_VERSION, + "rank": rank, + "capture_point": ( + "hash worker launched after checkpoint load and before the first P3O stats/train forward; " + "joined before optimizer prepare/update while parameters remain immutable" + ), + "parameter_sha256": aggregate.hexdigest(), + "tensor_count": len(tensors), + "total_numel": sum(int(record["numel"]) for record in tensors), + "tensors": tensors, + }, + ) + _STATE["parameters_captured"] = True + + +def _start_initial_parameter_capture(model: Any) -> None: + if _STATE["parameter_capture_thread"] is not None: + raise RuntimeError("Task40 Batch 7 parameter capture worker was already started") + + def worker() -> None: + try: + _capture_initial_parameters(model) + except BaseException as exc: # surfaced on the training thread by the join below + _STATE["parameter_capture_error"] = exc + + thread = threading.Thread(target=worker, name="task40-b7-parameter-sha256", daemon=False) + _STATE["parameter_capture_thread"] = thread + thread.start() + + +def _join_initial_parameter_capture() -> None: + thread = _STATE["parameter_capture_thread"] + if thread is None: + raise RuntimeError("Task40 Batch 7 parameter capture worker was not started") + thread.join() + error = _STATE["parameter_capture_error"] + if error is not None: + raise RuntimeError("Task40 Batch 7 initial parameter SHA-256 worker failed") from error + if not _STATE["parameters_captured"]: + raise RuntimeError("Task40 Batch 7 initial parameter SHA-256 worker produced no manifest") + + +def _capture_prepared_gradient_shards(model: Any, optimizer: Any) -> None: + import torch + + if _STATE["gradients_captured"]: + raise RuntimeError("Task40 Batch 7 attempted to capture gradients more than once") + output_dir: Path = _STATE["output_dir"] + rank = int(_STATE["rank"]) + gradient_dir = output_dir / "gradients" + gradient_dir.mkdir(exist_ok=True) + vector_dir = gradient_dir / f"rank{rank:05d}" + vector_dir.mkdir(exist_ok=False) + + aggregate = hashlib.sha256() + tensors = [] + missing_owned = [] + aggregate_l2_sq = 0.0 + aggregate_numel = 0 + names = {id(parameter): (name, parameter) for name, parameter in _named_parameters(model)} + inner_optimizers = getattr(optimizer, "chained_optimizers", [optimizer]) + shard_index = 0 + for optimizer_index, inner in enumerate(inner_optimizers): + if not hasattr(inner, "model_param_group_index_map") or not hasattr(inner, "_get_model_param_range_map"): + raise RuntimeError( + "Task40 Batch 7 gradient reconstruction requires Megatron DistributedOptimizer shard metadata" + ) + for model_parameter, (group_index, group_order) in inner.model_param_group_index_map.items(): + if id(model_parameter) not in names: + raise RuntimeError("Distributed optimizer owns a model parameter that is absent from named_parameters") + name, parameter = names[id(model_parameter)] + parameter_range = inner._get_model_param_range_map(model_parameter)["param"] + optimizer_parameter = inner.optimizer.param_groups[group_index]["params"][group_order] + gradient = optimizer_parameter.grad + record: dict[str, Any] = { + "index": shard_index, + "optimizer_index": optimizer_index, + "name": name, + "parameter_shape": list(parameter.shape), + "parameter_numel": parameter.numel(), + "range_start": int(parameter_range.start), + "range_end": int(parameter_range.end), + "requires_grad": bool(parameter.requires_grad), + "source": "prepared_distributed_optimizer_shard", + } + shard_index += 1 + if gradient is None: + record.update({"present": False, "numel": 0, "file": None}) + missing_owned.append(name) + tensors.append(record) + continue + + cpu, raw = _tensor_bytes(gradient) + if cpu.numel() != int(parameter_range.end) - int(parameter_range.start): + raise RuntimeError( + f"prepared gradient shard length mismatch for {name}: gradient={cpu.numel()}, " + f"range={parameter_range.start}:{parameter_range.end}" + ) + values = cpu.to(torch.float64) + l2_sq = float(torch.sum(values * values)) + finite = bool(torch.isfinite(values).all()) + file_name = f"shard_{record['index']:05d}.pt" + torch.save(cpu, vector_dir / file_name) + header = json.dumps( + { + "dtype": str(cpu.dtype), + "name": name, + "range_end": int(parameter_range.end), + "range_start": int(parameter_range.start), + "shape": list(cpu.shape), + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + aggregate.update(len(header).to_bytes(8, "little")) + aggregate.update(header) + aggregate.update(len(raw).to_bytes(8, "little")) + aggregate.update(raw) + aggregate_l2_sq += l2_sq + aggregate_numel += cpu.numel() + record.update( + { + "present": True, + "shape": list(cpu.shape), + "dtype": str(cpu.dtype), + "numel": cpu.numel(), + "finite": finite, + "l2_sq": l2_sq, + "l2_norm": l2_sq**0.5, + "max_abs": float(values.abs().max()) if values.numel() else 0.0, + "sum": float(values.sum()), + "sha256": hashlib.sha256(raw).hexdigest(), + "file": f"rank{rank:05d}/{file_name}", + } + ) + tensors.append(record) + + _write_json( + gradient_dir / f"summary_rank{rank}.json", + { + "format_version": FORMAT_VERSION, + "rank": rank, + "capture_point": "after DistributedOptimizer.prepare_grads and before clipping/update", + "gradient_sha256": aggregate.hexdigest(), + "shard_count": len(tensors), + "present_shard_count": sum(bool(record["present"]) for record in tensors), + "aggregate_numel": aggregate_numel, + "aggregate_l2_sq": aggregate_l2_sq, + "aggregate_l2_norm": aggregate_l2_sq**0.5, + "missing_owned_gradients": missing_owned, + "vector_files_saved": True, + "tensors": tensors, + }, + ) + _STATE["gradients_captured"] = True + + +def configure(args: Any) -> None: + """Install the Batch-7 BF16-only oracle observers.""" + import torch + import torch.distributed as dist + from megatron.core import mpu + + from relax.backends.megatron import p3o_step + + if _STATE["args"] is not None: + raise RuntimeError("Task40 Batch 7 revalidation configure() was called more than once") + if not bool(getattr(args, "bf16", False)) or bool(getattr(args, "fp16", False)): + raise ValueError("Task40 Batch 7 permits BF16 only") + if str(getattr(args, "qkv_format", "thd")) != "thd": + raise ValueError("Task40 Batch 7 requires THD") + if int(getattr(args, "micro_batch_size", 0)) != 1: + raise ValueError("Task40 Batch 7 requires micro-batch size 1") + if str(getattr(args, "p3o_ess_scope", "")) != "step": + raise ValueError("Task40 Batch 7 requires P3O step ESS scope") + fixture = Path(str(getattr(args, "load_debug_rollout_data", ""))) + if not fixture.is_file(): + raise ValueError(f"Task40 Step-0 fixture does not exist: {fixture}") + fixture_sha = _sha256(fixture) + if fixture_sha != EXPECTED_FIXTURE_SHA256: + raise ValueError( + f"Task40 Step-0 fixture SHA-256 mismatch: got {fixture_sha}, expected {EXPECTED_FIXTURE_SHA256}" + ) + + output_dir = Path(args.dump_details).parent / "oracle" + output_dir.mkdir(parents=True, exist_ok=True) + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + _STATE.update({"args": args, "output_dir": output_dir, "rank": rank}) + runtime = { + "format_version": FORMAT_VERSION, + "rank": rank, + "world_size": dist.get_world_size(), + "dp_rank": mpu.get_data_parallel_rank(with_context_parallel=False), + "dp_world_size": mpu.get_data_parallel_world_size(with_context_parallel=False), + "cp_rank": mpu.get_context_parallel_rank(), + "cp_world_size": mpu.get_context_parallel_world_size(), + "tp_rank": mpu.get_tensor_model_parallel_rank(), + "tp_world_size": mpu.get_tensor_model_parallel_world_size(), + "pp_rank": mpu.get_pipeline_model_parallel_rank(), + "pp_world_size": mpu.get_pipeline_model_parallel_world_size(), + "bf16": bool(args.bf16), + "fp16": bool(args.fp16), + "params_dtype": str(args.params_dtype), + "qkv_format": str(args.qkv_format), + "micro_batch_size": int(args.micro_batch_size), + "global_batch_size": int(args.global_batch_size), + "p3o_ess_scope": str(args.p3o_ess_scope), + "fixture": str(fixture), + "fixture_sha256": fixture_sha, + "expected_fixture_sha256": EXPECTED_FIXTURE_SHA256, + "parameter_capture_point": ( + "worker launched by before_train_step after checkpoint loading and before stats forward; " + "joined before optimizer prepare/update" + ), + "gradient_capture_point": "after DistributedOptimizer.prepare_grads and before clipping/update", + } + _write_json(output_dir / f"runtime_rank{rank}.json", runtime) + + original_local_stats = p3o_step._local_stats_from_batch + original_synchronize = p3o_step.synchronize_p3o_stats + + def observed_local_stats(local_args: Any, batch: dict[str, Any], log_probs: list[Any]) -> Any: + result = original_local_stats(local_args, batch, log_probs) + local_counter = int(_STATE["local_counter"]) + if not batch.get("__is_dummy__", False): + current = torch.cat(log_probs, dim=0) + behavior = torch.cat(batch["rollout_log_probs"], dim=0) + valid_mask = p3o_step.get_cp_local_valid_mask( + batch["total_lengths"], + batch["response_lengths"], + batch["loss_masks"], + local_args.qkv_format, + batch.get("max_seq_lens"), + batch.get("padded_total_lengths"), + dynamic_cp_size=batch.get("dynamic_cp_size"), + dynamic_cp_rank=batch.get("dynamic_cp_rank"), + ) + packed = batch.get("packed_seq_params") + total_lengths = [int(value) for value in batch["total_lengths"]] + artifact = { + "rank": rank, + "micro_batch_index": local_counter, + "current_log_probs": _cpu(current), + "rollout_log_probs": _cpu(behavior), + "valid_mask": _cpu(valid_mask.bool()), + "loss_masks": _cpu(batch["loss_masks"]), + "total_lengths": total_lengths, + "response_lengths": [int(value) for value in batch["response_lengths"]], + "token_ids": _cpu(batch.get("unconcat_tokens", batch.get("tokens"))), + "position_ids": [torch.arange(length, dtype=torch.int64) for length in total_lengths], + "max_seq_lens": _cpu(batch.get("max_seq_lens")), + "padded_total_lengths": _cpu(batch.get("padded_total_lengths")), + "cu_seqlens_q": _cpu(getattr(packed, "cu_seqlens_q", None)), + "cu_seqlens_kv": _cpu(getattr(packed, "cu_seqlens_kv", None)), + "max_seqlen_q": getattr(packed, "max_seqlen_q", None), + "max_seqlen_kv": getattr(packed, "max_seqlen_kv", None), + "relax_total_lengths": _cpu(getattr(packed, "_relax_total_lengths", None)), + "relax_attention_pad_multiple": getattr(packed, "_relax_attention_pad_multiple", None), + "relax_cu_seqlens_cpu": _cpu(getattr(packed, "_relax_cu_seqlens_cpu", None)), + "local_s1": _cpu(result[0].sum_ratio), + "local_s2": _cpu(result[0].sum_ratio_sq), + "local_n": _cpu(result[0].valid_token_count), + "invalid_flag": _cpu(result[1]), + } + torch.save(artifact, output_dir / f"vectors_rank{rank}_micro{local_counter}.pt") + _STATE["local_counter"] = local_counter + 1 + return result + + def observed_synchronize(*sync_args: Any, **sync_kwargs: Any) -> Any: + reduced = original_synchronize(*sync_args, **sync_kwargs) + sync_counter = int(_STATE["sync_counter"]) + torch.save( + { + "rank": rank, + "sync_index": sync_counter, + "s1": _cpu(reduced.sum_ratio), + "s2": _cpu(reduced.sum_ratio_sq), + "n": _cpu(reduced.valid_token_count), + }, + output_dir / f"global_stats_rank{rank}_sync{sync_counter}.pt", + ) + _STATE["sync_counter"] = sync_counter + 1 + return reduced + + p3o_step._local_stats_from_batch = observed_local_stats + p3o_step.synchronize_p3o_stats = observed_synchronize + + +def before_train_step( + args: Any, + rollout_id: int, + step_id: int, + model: Any, + optimizer: Any, + opt_param_scheduler: Any, +) -> None: + """Capture parameters and wrap the one optimizer step for gradients.""" + del args, opt_param_scheduler + if rollout_id != 0 or step_id != 0: + raise RuntimeError(f"Task40 Batch 7 permits only rollout=0/step=0, got {rollout_id=}, {step_id=}") + _start_initial_parameter_capture(model) + optimizer_id = id(optimizer) + if optimizer_id in _STATE["optimizer_ids"]: + raise RuntimeError("Task40 Batch 7 optimizer step observer was already installed") + _STATE["optimizer_ids"].add(optimizer_id) + original_prepare_grads = optimizer.prepare_grads + + def observed_prepare_grads(*prepare_args: Any, **prepare_kwargs: Any) -> Any: + _join_initial_parameter_capture() + found_inf = original_prepare_grads(*prepare_args, **prepare_kwargs) + if not found_inf: + _capture_prepared_gradient_shards(model, optimizer) + return found_inf + + optimizer.prepare_grads = observed_prepare_grads diff --git a/tests/scripts/local_p3o_task40/test_step0_revalidation.py b/tests/scripts/local_p3o_task40/test_step0_revalidation.py new file mode 100644 index 000000000..a9a235c35 --- /dev/null +++ b/tests/scripts/local_p3o_task40/test_step0_revalidation.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import sys +from pathlib import Path + +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from scripts.local_p3o_task40 import analyze_step0_revalidation as analysis # noqa: E402 +from scripts.local_p3o_task40 import step0_revalidation as hook # noqa: E402 + + +def test_cu_seqlens_derivation_makes_cp2_relation_explicit() -> None: + assert analysis._canonical_cu([282], 128) == [0, 282, 384] + assert analysis._expected_source_cu([282], 1, 128) == [0, 282, 384] + assert analysis._expected_source_cu([282], 2, 128) == [0, 284, 512] + assert analysis._expected_source_cu([219], 1, 128) == [0, 219, 256] + + +def test_numeric_comparison_uses_frozen_relative_and_near_zero_gates() -> None: + assert analysis._numeric_comparison(2.0, 2.0 + 1e-6)["pass"] + assert not analysis._numeric_comparison(2.0, 2.0 + 3e-6)["pass"] + assert analysis._numeric_comparison(0.0, 1e-9)["pass"] + assert not analysis._numeric_comparison(0.0, 2e-9)["pass"] + + +def test_tensor_bytes_hashes_bf16_losslessly() -> None: + tensor = torch.tensor([1.0, -2.0], dtype=torch.bfloat16) + cpu, raw = hook._tensor_bytes(tensor) + assert torch.equal(cpu, tensor) + assert len(raw) == tensor.numel() * tensor.element_size() + + +def test_parameter_hash_worker_is_joined_before_update(tmp_path: Path, monkeypatch) -> None: + state = { + **hook._STATE, + "output_dir": tmp_path, + "rank": 0, + "parameters_captured": False, + "parameter_capture_error": None, + "parameter_capture_thread": None, + } + monkeypatch.setattr(hook, "_STATE", state) + hook._start_initial_parameter_capture([torch.nn.Linear(2, 2)]) + hook._join_initial_parameter_capture() + assert state["parameters_captured"] + assert (tmp_path / "initial_parameters_rank0.json").is_file() + + +def test_gradient_comparator_reports_full_vector_rel_l2_and_cosine(tmp_path: Path) -> None: + summaries = [] + for label, vector in (("reference", torch.tensor([1.0, 2.0])), ("candidate", torch.tensor([1.0, 2.0]))): + run_dir = tmp_path / label + vector_dir = run_dir / "oracle" / "gradients" / "rank00000" + vector_dir.mkdir(parents=True) + torch.save(vector, vector_dir / "shard_00000.pt") + summary = { + "rank": 0, + "tensors": [ + { + "name": "chunk000.weight", + "present": True, + "file": "rank00000/shard_00000.pt", + "parameter_numel": 2, + "range_start": 0, + "range_end": 2, + } + ], + } + (run_dir / "oracle" / "gradients" / "summary_rank0.json").write_text(__import__("json").dumps(summary)) + summaries.append(run_dir) + comparison = analysis._compare_gradients(*summaries) + assert comparison["pass"] + assert comparison["full_parameter_relative_l2"] == 0.0 + assert comparison["full_parameter_cosine"] == 1.0 + + +def test_reconstruct_gradient_joins_distributed_optimizer_shards(tmp_path: Path) -> None: + records = [] + for rank, (start, values) in enumerate(((0, torch.tensor([1.0, 2.0])), (2, torch.tensor([3.0, 4.0])))): + path = tmp_path / f"rank{rank}.pt" + torch.save(values, path) + records.append( + { + "name": "chunk000.weight", + "parameter_numel": 4, + "range_start": start, + "range_end": start + values.numel(), + "path": path, + } + ) + assert torch.equal(analysis._reconstruct_gradient(records), torch.tensor([1.0, 2.0, 3.0, 4.0])) + + +def test_launcher_is_frozen_to_four_bf16_step_scope_cells() -> None: + launcher = (REPO_ROOT / "scripts/local_p3o_task40/run_step0_revalidation.sh").read_text() + for topology in ("dp1", "dp4cp1", "dp2cp1", "dp2cp2"): + assert topology in launcher + assert "P3O_ESS_SCOPE=step" in launcher + assert "P3O_MICRO_BATCH_SIZE=1" in launcher + assert "step0_revalidation.configure" in launcher + assert "fp32" not in launcher.lower() From ef854a3cfb3eb692e083d14e12fbd96a8b1c7af5 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:03:08 +0800 Subject: [PATCH 18/30] fix(p3o): enforce deterministic partition kernels --- examples/algorithms/p3o/README.md | 27 ++-- examples/algorithms/p3o/README_zh.md | 23 ++- examples/algorithms/p3o/common_a100x4.sh | 10 ++ relax/backends/megatron/data.py | 6 +- relax/backends/megatron/initialize.py | 21 +++ relax/backends/megatron/model.py | 65 +++++---- relax/backends/megatron/model_provider.py | 2 + .../megatron/test_model_provider_vpp.py | 14 ++ ...test_p3o_attention_partition_invariance.py | 137 +++++++++++++----- .../backends/megatron/test_p3o_initialize.py | 97 +++++++++++++ tests/examples/algorithms/p3o/test_configs.py | 5 + 11 files changed, 329 insertions(+), 78 deletions(-) create mode 100644 tests/backends/megatron/test_p3o_initialize.py diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md index 1da991d64..e4206239a 100644 --- a/examples/algorithms/p3o/README.md +++ b/examples/algorithms/p3o/README.md @@ -87,16 +87,23 @@ reward would define an unvalidated hybrid objective. The reward/verifier name valid for compatible datasets. With `--context-parallel-size > 1`, P3O automatically runs every THD -self-attention layer 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 the QKV projection and attention kernel see the -same token order and shape as CP1. It duplicates full-sequence QKV, attention, -and projection 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 validated 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. +TransformerLayer 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 QKV, attention, residual, and MLP kernels see the same +token order and shape as CP1. It duplicates full-layer 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. +Batch-invariant kernels cover the final normalization and LM head outside the +wrapped TransformerLayers; P3O fails closed if either mode is absent. 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 diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md index d8b6693b2..fb681cf7a 100644 --- a/examples/algorithms/p3o/README_zh.md +++ b/examples/algorithms/p3o/README_zh.md @@ -70,13 +70,22 @@ OPD advantage replacement 或 OPD-only reward 都会形成未经验证的混合 reward/verifier 名称 `P3O_RM_TYPE=mopd` 与 `--use-opd` 训练功能无关，兼容数据集 仍可使用。 -当 `--context-parallel-size > 1` 时，P3O 会自动在每层 THD self-attention 之前重建 -全序列，完成计算后再把输出切回各 CP rank。该严格路径对 `micro-batch` 和 `step` -两种 ESS scope 都生效，使 QKV 投影和 attention kernel 看到与 CP1 相同的 token -顺序和形状。它会在每个 CP rank 上重复完整序列的 QKV、attention 和输出投影计算及 -activation，因此峰值显存和计算量都高于原生 CP；长上下文建议开启整层 activation -recomputation。当前已验证合同是标准 zig-zag THD 且 tensor parallel size 为 1。如果全序列路径发生 -OOM，P3O 会终止并拒绝 CP>1，不会静默回退到数值不等价的原生 CP kernel order。 +当 `--context-parallel-size > 1` 时，P3O 会自动在每个 THD TransformerLayer 之前 +重建全序列，完成整层计算后再把输出切回各 CP rank。该严格路径对 `micro-batch` 和 +`step` 两种 ESS scope 都生效，使 QKV、attention、残差和 MLP kernel 看到与 CP1 +相同的 token 顺序和形状。它会在每个 CP rank 上重复完整层的计算和 activation， +因此峰值显存和计算量都高于原生 CP；长上下文建议开启整层 activation +recomputation。当前支持的合同是标准 zig-zag THD 且 tensor parallel size 为 1。 +如果全序列路径发生 OOM，P3O 会终止并拒绝 CP>1，不会静默回退到数值不等价的 +原生 CP kernel order。 + +严格分区不变性还要求 Megatron `--deterministic-mode` 和 +`--batch-invariant-mode`。示例脚本会为配对的 P3O/GRPO run 同时启用两者，并在模型 +构建前向 Ray runtime 注入 `NCCL_ALGO=Ring`、 +`NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` 和 +`CUBLAS_WORKSPACE_CONFIG=:4096:8`。batch-invariant kernel 覆盖 wrapped +TransformerLayer 之外的最终 normalization 与 LM head；缺少任一模式时 P3O 会 +fail closed。 正式默认值为 G=16、global batch 64、micro-batch 1、rollout batch 4、response length 4096 和 30 个 optimizer step（`--num-rollout 30`）。计划配对 seed 为 42、 diff --git a/examples/algorithms/p3o/common_a100x4.sh b/examples/algorithms/p3o/common_a100x4.sh index 21a62649f..fffa204d7 100755 --- a/examples/algorithms/p3o/common_a100x4.sh +++ b/examples/algorithms/p3o/common_a100x4.sh @@ -294,6 +294,8 @@ P3O_build_args() { --rollout-seed "${P3O_SEED}" --attention-dropout 0.0 --hidden-dropout 0.0 + --deterministic-mode + --batch-invariant-mode --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 --attention-backend flash @@ -393,6 +395,11 @@ P3O_run() { echo "pipeline_model_parallel_size=${P3O_PIPELINE_MODEL_PARALLEL_SIZE}" echo "nccl_debug=${P3O_NCCL_DEBUG}" echo "torch_distributed_debug=${P3O_TORCH_DISTRIBUTED_DEBUG}" + echo "deterministic_mode=1" + echo "batch_invariant_mode=1" + echo "nccl_algo=Ring" + echo "nvte_allow_nondeterministic_algo=0" + echo "cublas_workspace_config=:4096:8" echo "behavior_temperature=${P3O_BEHAVIOR_TEMPERATURE}" echo "ray_job_id=${P3O_JOB_ID}" echo "repo=${P3O_REPO_ROOT}" @@ -446,6 +453,9 @@ env_vars = { "OPENBLAS_NUM_THREADS": os.environ["P3O_RUNTIME_OPENBLAS_NUM_THREADS"], "NCCL_NVLS_ENABLE": os.environ["P3O_RUNTIME_NCCL_NVLS_ENABLE"], "NVSHMEM_DISABLE_NCCL": os.environ["P3O_RUNTIME_NVSHMEM_DISABLE_NCCL"], + "NCCL_ALGO": "Ring", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", } if os.environ["P3O_RUNTIME_CLEAR_PROXIES"] == "1": diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index a8d5f8323..508d5075c 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -482,9 +482,9 @@ def get_batch( packed_seq_params.local_cp_size = cp_size packed_seq_params.cp_group = cp_group if getattr(get_args(), "advantage_estimator", None) == "p3o": - # P3O's strict CP attention reconstructs the CP1 QKV shape without - # changing this existing token layout. Keep the host boundaries from - # construction so every layer avoids a device-to-host synchronization. + # P3O's strict CP path reconstructs the CP1 TransformerLayer shape + # without changing this existing token layout. Keep the host + # boundaries so every layer avoids a device-to-host synchronization. packed_seq_params._relax_total_lengths = list(batch["total_lengths"]) packed_seq_params._relax_attention_pad_multiple = pad_size packed_seq_params._relax_cu_seqlens_cpu = cu_seqlens_cpu diff --git a/relax/backends/megatron/initialize.py b/relax/backends/megatron/initialize.py index 7d763a19a..175ab61f1 100644 --- a/relax/backends/megatron/initialize.py +++ b/relax/backends/megatron/initialize.py @@ -1,5 +1,6 @@ import random import socket +from argparse import Namespace import numpy as np import torch @@ -7,6 +8,7 @@ from megatron.core import mpu, tensor_parallel from megatron.core.config import set_experimental_flag from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator +from megatron.core.transformer.custom_layers.batch_invariant_kernels import enable_batch_invariant_mode from megatron.training.global_vars import _build_tokenizer, set_args from relax.utils.logging_utils import get_logger @@ -15,6 +17,24 @@ logger = get_logger(__name__) +def _configure_p3o_partition_invariance(args: Namespace) -> None: + """Enable batch-invariant kernels and fail closed for incomplete P3O + mode.""" + batch_invariant_mode = getattr(args, "batch_invariant_mode", False) + if getattr(args, "advantage_estimator", None) == "p3o": + missing: list[str] = [] + if not getattr(args, "deterministic_mode", False): + missing.append("--deterministic-mode") + if not batch_invariant_mode: + missing.append("--batch-invariant-mode") + if missing: + raise RuntimeError( + "P3O strict partition invariance requires " + " and ".join(missing) + "; refusing a non-equivalent run" + ) + if batch_invariant_mode: + enable_batch_invariant_mode() + + def _set_random_seed( seed_: int, data_parallel_random_init: bool = False, @@ -76,6 +96,7 @@ def _initialize_distributed(args, get_embedding_ranks=None, get_position_embeddi def init(args): set_args(args) + _configure_p3o_partition_invariance(args) if getattr(args, "disable_jit_fuser", False): from megatron.core.jit import disable_jit_fuser diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index f2e8e3f6d..ee81eb18d 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -57,7 +57,7 @@ P3O_CP_ATTENTION_OOM_ERROR = ( - "P3O strict mode refuses context-parallel size greater than one after full-sequence attention OOM" + "P3O strict mode refuses context-parallel size greater than one after full-sequence TransformerLayer OOM" ) @@ -191,45 +191,56 @@ def _restore_p3o_attention_output_layout( return output -def _build_p3o_full_sequence_attention_forward(attention: torch.nn.Module) -> Callable: - """Wrap one SelfAttention instance with P3O's strict CP-equivalent path.""" - original_forward = attention.forward +def _build_p3o_full_sequence_layer_forward(layer: torch.nn.Module) -> Callable: + """Wrap one TransformerLayer with P3O's strict CP-equivalent path.""" + original_forward = layer.forward @wraps(original_forward) def full_sequence_forward( hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None, - key_value_states: torch.Tensor | None = None, - inference_context: object | None = None, + attention_mask: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, rotary_pos_emb: object | None = None, rotary_pos_cos: torch.Tensor | None = None, rotary_pos_sin: torch.Tensor | None = None, rotary_pos_cos_sin: torch.Tensor | None = None, attention_bias: torch.Tensor | None = None, + inference_context: object | None = None, packed_seq_params: object | None = None, - sequence_len_offset: int | None = None, + sequence_len_offset: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, *, inference_params: object | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + **kwargs: object, + ) -> tuple[torch.Tensor, object | None]: if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": - raise RuntimeError("P3O full-sequence attention requires THD packed_seq_params when CP>1") + raise RuntimeError("P3O full-sequence layers require THD packed_seq_params when CP>1") + attention = getattr(layer, "self_attention", None) + if attention is None: + raise RuntimeError("P3O full-sequence layer requires a self_attention module") cp_size, cp_group, cp_rank = _resolve_p3o_attention_cp(attention, packed_seq_params) if cp_size == 1: return original_forward( hidden_states, attention_mask, - key_value_states, - inference_context, + context, + context_mask, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, rotary_pos_cos_sin, attention_bias, + inference_context, packed_seq_params, sequence_len_offset, + padding_mask, inference_params=inference_params, + **kwargs, ) + if padding_mask is not None: + raise RuntimeError("P3O strict full-sequence layers do not support a CP-local padding_mask") cu_seqlens_cpu = getattr(packed_seq_params, "_relax_cu_seqlens_cpu", None) if cu_seqlens_cpu is None: @@ -262,26 +273,29 @@ def full_sequence_forward( if hasattr(full_packed_seq_params, "cu_seqlens_q_padded"): full_packed_seq_params.cu_seqlens_q_padded = None full_packed_seq_params.cu_seqlens_kv_padded = None - full_output, bias = original_forward( + full_output, full_context = original_forward( canonical_hidden_states, attention_mask, - key_value_states, - inference_context, + context, + context_mask, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, rotary_pos_cos_sin, attention_bias, + inference_context, full_packed_seq_params, sequence_len_offset, + padding_mask, inference_params=inference_params, + **kwargs, ) source_layout_output = _restore_p3o_attention_output_layout( full_output, cu_seqlens_cpu, total_lengths, ) - return _p3o_cp_slice(source_layout_output, cu_seqlens_cpu, cp_size, cp_rank), bias + return _p3o_cp_slice(source_layout_output, cu_seqlens_cpu, cp_size, cp_rank), full_context except torch.cuda.OutOfMemoryError as exc: raise RuntimeError(P3O_CP_ATTENTION_OOM_ERROR) from exc finally: @@ -298,14 +312,15 @@ def _install_p3o_full_sequence_attention(args: Namespace, model: torch.nn.Module installed = 0 for module in model.modules(): - if type(module).__name__ != "SelfAttention" or getattr( - module, "_relax_p3o_full_sequence_attention_installed", False + if type(module).__name__ != "TransformerLayer" or getattr( + module, "_relax_p3o_full_sequence_layer_installed", False ): continue - if not hasattr(module, "pg_collection") or not hasattr(module.pg_collection, "cp"): - raise RuntimeError("P3O full-sequence attention requires SelfAttention.pg_collection.cp") - module.forward = _build_p3o_full_sequence_attention_forward(module) - module._relax_p3o_full_sequence_attention_installed = True + attention = getattr(module, "self_attention", None) + if attention is None or not hasattr(attention, "pg_collection") or not hasattr(attention.pg_collection, "cp"): + raise RuntimeError("P3O full-sequence layer requires SelfAttention.pg_collection.cp") + module.forward = _build_p3o_full_sequence_layer_forward(module) + module._relax_p3o_full_sequence_layer_installed = True installed += 1 return installed @@ -681,11 +696,11 @@ def setup_model_and_optimizer( ) if getattr(args, "advantage_estimator", None) == "p3o" and getattr(args, "context_parallel_size", 1) > 1: - installed_attention_modules = sum( + installed_full_sequence_layers = sum( _install_p3o_full_sequence_attention(args, model_chunk) for model_chunk in model ) - if installed_attention_modules == 0: - raise RuntimeError("P3O context parallelism requires at least one supported SelfAttention module") + if installed_full_sequence_layers == 0: + raise RuntimeError("P3O context parallelism requires at least one supported TransformerLayer") # Some model providers (e.g., Qwen3VLGPTModel) rebuild the decoder in __init__, # which causes duplicate RoutingReplay registrations. Rebuild the list from diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 0893b3366..0331f4b52 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -212,6 +212,8 @@ def wrapped_model_provider( # Override provider attributes with matching args values bridge_keys = [ "attention_backend", + "deterministic_mode", + "batch_invariant_mode", "tensor_model_parallel_size", "sequence_parallel", "pipeline_model_parallel_size", diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 9a273ebec..e45ec03e8 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -29,6 +29,8 @@ def __init__(self, model=None): self.bf16 = False self.params_dtype = None self.vision_dp_when_cp = False + self.deterministic_mode = False + self.batch_invariant_mode = False def finalize(self): self.finalized = True @@ -133,6 +135,8 @@ def _bridge_args(**overrides): "sequence_parallel": True, "pipeline_model_parallel_size": 4, "virtual_pipeline_model_parallel_size": 2, + "deterministic_mode": True, + "batch_invariant_mode": True, "context_parallel_size": 1, "expert_model_parallel_size": 1, "expert_tensor_parallel_size": 1, @@ -183,6 +187,16 @@ def test_bridge_provider_receives_virtual_pipeline_size(monkeypatch): assert provider.calls == [{"pre_process": True, "post_process": False, "vp_stage": 1}] +def test_bridge_provider_receives_partition_invariance_modes(monkeypatch): + module, provider = _load_model_provider(monkeypatch) + + module.get_model_provider_func(_bridge_args(), role="actor") + + assert provider.deterministic_mode is True + assert provider.batch_invariant_mode is True + assert provider.finalized + + def test_bridge_provider_uses_fp32_when_precision_flags_are_disabled(monkeypatch): module, provider = _load_model_provider(monkeypatch) diff --git a/tests/backends/megatron/test_p3o_attention_partition_invariance.py b/tests/backends/megatron/test_p3o_attention_partition_invariance.py index bb7bf277a..04526c874 100644 --- a/tests/backends/megatron/test_p3o_attention_partition_invariance.py +++ b/tests/backends/megatron/test_p3o_attention_partition_invariance.py @@ -21,10 +21,10 @@ ) -class _AttentionRoot(torch.nn.Module): - def __init__(self, attention: torch.nn.Module): +class _LayerRoot(torch.nn.Module): + def __init__(self, layer: torch.nn.Module): super().__init__() - self.attention = attention + self.layer = layer class SelfAttention(torch.nn.Module): @@ -77,6 +77,58 @@ def forward( return output, torch.zeros(projected.shape[-1], dtype=projected.dtype) +class TransformerLayer(torch.nn.Module): + """Megatron-shaped layer with a token-shape-sensitive post-attention + MLP.""" + + def __init__(self, cp_group: object, attention_weight: torch.Tensor, mlp_weight: torch.Tensor): + super().__init__() + self.self_attention = SelfAttention(cp_group, attention_weight) + self.mlp_weight = torch.nn.Parameter(mlp_weight.clone()) + self.seen_token_counts: list[int] = [] + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_cos: torch.Tensor | None = None, + rotary_pos_sin: torch.Tensor | None = None, + rotary_pos_cos_sin: torch.Tensor | None = None, + attention_bias: torch.Tensor | None = None, + inference_context: object | None = None, + packed_seq_params: object | None = None, + sequence_len_offset: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + *, + inference_params: object | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del context_mask, padding_mask + attention_output, _ = self.self_attention( + hidden_states, + attention_mask, + None, + inference_context, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + rotary_pos_cos_sin, + attention_bias, + packed_seq_params, + sequence_len_offset, + inference_params=inference_params, + ) + token_count = attention_output.shape[0] + self.seen_token_counts.append(token_count) + # Real fused MLP/residual kernels are numerically shape-sensitive. Make + # that boundary explicit so an attention-only adapter cannot satisfy + # this whole-layer partition-invariance test accidentally. + output = (attention_output @ self.mlp_weight) * (1.0 + token_count / 100.0) + return output, context + + def _free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) @@ -111,24 +163,31 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non ] full_hidden = torch.cat(cp_padded_samples) canonical_hidden = torch.cat([*real_samples, torch.zeros(6, 1, hidden_size, dtype=torch.float64)]) - weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) + attention_weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) + mlp_weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) reference_hidden = canonical_hidden.clone().requires_grad_(True) - reference_weight = weight.clone().requires_grad_(True) - reference_projected = reference_hidden @ reference_weight - reference_output = torch.cat( - [ - torch.cumsum(reference_projected[start:end], dim=0) - for start, end in zip(canonical_cu_seqlens[:-1], canonical_cu_seqlens[1:], strict=True) - ] + reference_layer = TransformerLayer( + SimpleNamespace(size=lambda: 1, rank=lambda: 0), attention_weight, mlp_weight + ) + reference_packed_seq_params = SimpleNamespace( + qkv_format="thd", + cu_seqlens_q=torch.tensor(canonical_cu_seqlens, dtype=torch.int32), + cu_seqlens_kv=torch.tensor(canonical_cu_seqlens, dtype=torch.int32), + local_cp_size=1, + cp_group=SimpleNamespace(size=lambda: 1, rank=lambda: 0), + ) + reference_output, _ = reference_layer( + reference_hidden, + packed_seq_params=reference_packed_seq_params, ) reference_mask = torch.zeros(len(reference_output), 1, 1, dtype=torch.float64) reference_mask[: sum(total_lengths)] = 1.0 (reference_output * reference_mask).square().sum().backward() local_hidden = gdn_cp_slice(full_hidden, cu_seqlens, world_size, rank).clone().requires_grad_(True) - attention = SelfAttention(dist.group.WORLD, weight) - root = _AttentionRoot(attention) + layer = TransformerLayer(dist.group.WORLD, attention_weight, mlp_weight) + root = _LayerRoot(layer) assert _install_p3o_full_sequence_attention(_p3o_args(scope), root) == 1 packed_seq_params = SimpleNamespace( @@ -141,7 +200,7 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non _relax_attention_pad_multiple=8, _relax_cu_seqlens_cpu=cu_seqlens, ) - output, bias = attention(local_hidden, None, packed_seq_params=packed_seq_params) + output, context = layer(local_hidden, packed_seq_params=packed_seq_params) expected_full_output = torch.zeros_like(full_hidden) source_offset = 0 canonical_offset = 0 @@ -153,8 +212,9 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non canonical_offset += total_length expected_output = gdn_cp_slice(expected_full_output, cu_seqlens, world_size, rank) assert torch.equal(output.detach(), expected_output) - assert torch.equal(bias, torch.zeros(hidden_size, dtype=torch.float64)) - assert attention.seen_local_cp_sizes == [1] + assert context is None + assert layer.self_attention.seen_local_cp_sizes == [1] + assert layer.seen_token_counts == [len(canonical_hidden)] full_loss_mask = torch.zeros(len(full_hidden), 1, 1, dtype=torch.float64) source_offset = 0 @@ -174,8 +234,9 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non canonical_offset += total_length expected_input_grad = gdn_cp_slice(expected_full_input_grad, cu_seqlens, world_size, rank) torch.testing.assert_close(local_hidden.grad, expected_input_grad, atol=1e-10, rtol=1e-10) - dist.all_reduce(attention.weight.grad, op=dist.ReduceOp.SUM) - torch.testing.assert_close(attention.weight.grad, reference_weight.grad, atol=1e-10, rtol=1e-10) + for parameter, reference_parameter in zip(layer.parameters(), reference_layer.parameters(), strict=True): + dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) + torch.testing.assert_close(parameter.grad, reference_parameter.grad, atol=1e-10, rtol=1e-10) finally: dist.destroy_process_group() @@ -186,46 +247,56 @@ def test_p3o_full_sequence_attention_cpu_cp2_matches_cp1(scope: str) -> None: def test_p3o_full_sequence_attention_gate_is_p3o_cp_only() -> None: - attention = SelfAttention(SimpleNamespace(size=lambda: 1, rank=lambda: 0), torch.eye(2)) - root = _AttentionRoot(attention) - original_forward = attention.forward + layer = TransformerLayer( + SimpleNamespace(size=lambda: 1, rank=lambda: 0), + torch.eye(2), + torch.eye(2), + ) + root = _LayerRoot(layer) + original_forward = layer.forward assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch", context_parallel_size=1), root) == 0 - assert attention.forward == original_forward + assert layer.forward == original_forward assert ( _install_p3o_full_sequence_attention( _p3o_args("micro-batch", advantage_estimator="grpo", context_parallel_size=2), root ) == 0 ) - assert attention.forward == original_forward + assert layer.forward == original_forward assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 1 - installed_forward = attention.forward + installed_forward = layer.forward assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 0 - assert attention.forward == installed_forward + assert layer.forward == installed_forward def test_p3o_full_sequence_attention_oom_refuses_native_cp_fallback(monkeypatch: pytest.MonkeyPatch) -> None: from relax.backends.megatron import model as model_module cp_group = SimpleNamespace(size=lambda: 2, rank=lambda: 0) - attention = SelfAttention(cp_group, torch.eye(2)) - root = _AttentionRoot(attention) - assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 1 + layer = TransformerLayer(cp_group, torch.eye(2), torch.eye(2)) + leaked_group = SimpleNamespace(size=lambda: 1, rank=lambda: 0) - def raise_oom(*args: object, **kwargs: object) -> torch.Tensor: + def raise_oom_after_group_switch(*args: object, **kwargs: object) -> tuple[torch.Tensor, None]: del args, kwargs + layer.self_attention.pg_collection.cp = leaked_group raise torch.cuda.OutOfMemoryError("synthetic OOM") - monkeypatch.setattr(model_module, "_p3o_cp_gather_full", raise_oom) + layer.forward = raise_oom_after_group_switch + root = _LayerRoot(layer) + assert _install_p3o_full_sequence_attention(_p3o_args("micro-batch"), root) == 1 + + monkeypatch.setattr(model_module, "_p3o_cp_gather_full", lambda hidden, *args: torch.cat([hidden, hidden])) packed_seq_params = SimpleNamespace( qkv_format="thd", cu_seqlens_q=torch.tensor([0, 8], dtype=torch.int32), cu_seqlens_kv=torch.tensor([0, 8], dtype=torch.int32), local_cp_size=None, cp_group=None, + _relax_total_lengths=[5], + _relax_attention_pad_multiple=8, + _relax_cu_seqlens_cpu=[0, 8], ) with pytest.raises(RuntimeError, match=P3O_CP_ATTENTION_OOM_ERROR): - attention(torch.randn(4, 1, 2), None, packed_seq_params=packed_seq_params) - assert attention.seen_local_cp_sizes == [] - assert attention.pg_collection.cp is cp_group + layer(torch.randn(4, 1, 2), packed_seq_params=packed_seq_params) + assert layer.self_attention.pg_collection.cp is cp_group diff --git a/tests/backends/megatron/test_p3o_initialize.py b/tests/backends/megatron/test_p3o_initialize.py new file mode 100644 index 000000000..3a61c74aa --- /dev/null +++ b/tests/backends/megatron/test_p3o_initialize.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from argparse import Namespace + +import pytest +import torch + +from relax.backends.megatron import initialize + + +def _args(**overrides: object) -> Namespace: + values = { + "advantage_estimator": "p3o", + "batch_invariant_mode": True, + "deterministic_mode": True, + } + values.update(overrides) + return Namespace(**values) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"deterministic_mode": False}, "deterministic-mode"), + ({"batch_invariant_mode": False}, "batch-invariant-mode"), + ], +) +def test_p3o_partition_modes_fail_closed(overrides: dict[str, object], message: str) -> None: + with pytest.raises(RuntimeError, match=message): + initialize._configure_p3o_partition_invariance(_args(**overrides)) + + +def test_p3o_partition_modes_enable_batch_invariant_kernels(monkeypatch: pytest.MonkeyPatch) -> None: + calls = [] + monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: calls.append("enabled")) + + initialize._configure_p3o_partition_invariance(_args()) + + assert calls == ["enabled"] + + +def test_non_p3o_does_not_change_partition_modes(monkeypatch: pytest.MonkeyPatch) -> None: + calls = [] + monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: calls.append("enabled")) + + initialize._configure_p3o_partition_invariance( + _args(advantage_estimator="grpo", batch_invariant_mode=False, deterministic_mode=False) + ) + + assert calls == [] + + +def test_non_p3o_honors_explicit_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch) -> None: + calls = [] + monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: calls.append("enabled")) + + initialize._configure_p3o_partition_invariance( + _args(advantage_estimator="grpo", batch_invariant_mode=True, deterministic_mode=True) + ) + + assert calls == ["enabled"] + + +def _logical_dp_gradient(data_parallel_size: int) -> torch.Tensor: + torch.manual_seed(11) + weight = torch.randn(7, 5, dtype=torch.bfloat16) + inputs = [torch.randn(length, 5, dtype=torch.bfloat16) for length in (3, 5, 4, 6, 2, 7, 4, 5)] + targets = [torch.randn(value.shape[0], 7, dtype=torch.bfloat16) for value in inputs] + rank_gradients = [] + for indices in torch.tensor_split(torch.arange(len(inputs)), data_parallel_size): + main_grad = torch.zeros_like(weight, dtype=torch.float32) + for index in indices.tolist(): + parameter = weight.clone().requires_grad_(True) + output = torch.nn.functional.linear(inputs[index], parameter) + ((output - targets[index]) ** 2).sum().backward() + main_grad += parameter.grad.float() + rank_gradients.append(main_grad) + return sum(rank_gradients) + + +def test_p3o_bf16_mbs1_parameter_gradient_matches_across_dp_partitions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: None) + initialize._configure_p3o_partition_invariance(_args()) + + reference = _logical_dp_gradient(1) + for data_parallel_size in (2, 4): + candidate = _logical_dp_gradient(data_parallel_size) + reference64 = reference.double().flatten() + candidate64 = candidate.double().flatten() + relative_l2 = torch.linalg.vector_norm(candidate64 - reference64) / torch.linalg.vector_norm(reference64) + cosine = torch.dot(reference64, candidate64) / ( + torch.linalg.vector_norm(reference64) * torch.linalg.vector_norm(candidate64) + ) + assert relative_l2 <= 1e-6 + assert cosine >= 1 - 1e-9 diff --git a/tests/examples/algorithms/p3o/test_configs.py b/tests/examples/algorithms/p3o/test_configs.py index dc6f32704..06e35525b 100644 --- a/tests/examples/algorithms/p3o/test_configs.py +++ b/tests/examples/algorithms/p3o/test_configs.py @@ -273,6 +273,8 @@ def test_p3o_configs_freeze_required_formal_values(): assert _option_value(args, "--weight-decay") == "0.01" assert "--calculate-per-token-loss" in args assert "--use-rollout-logprobs" in args + assert "--deterministic-mode" in args + assert "--batch-invariant-mode" in args assert "--rollout-shuffle" in args assert "--colocate" in args assert "--fully-async" not in args @@ -505,6 +507,9 @@ def test_p3o_runner_executes_all_scenarios_with_fake_ray( assert runtime_env["NCCL_DEBUG"] == "WARN" assert runtime_env["TORCH_DISTRIBUTED_DEBUG"] == "OFF" assert runtime_env["RAY_OVERRIDE_JOB_RUNTIME_ENV"] == "1" + assert runtime_env["NCCL_ALGO"] == "Ring" + assert runtime_env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" + assert runtime_env["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" for proxy_name in ( "HTTP_PROXY", "HTTPS_PROXY", From 91e731694aa4d1be93d6e71b171cb2cb49fc948f Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:56:57 +0800 Subject: [PATCH 19/30] fix(p3o): pin retry verdict provenance --- scripts/local_p3o_task40/analyze_step0_revalidation.py | 3 ++- tests/scripts/local_p3o_task40/test_step0_revalidation.py | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/scripts/local_p3o_task40/analyze_step0_revalidation.py b/scripts/local_p3o_task40/analyze_step0_revalidation.py index 01ab821ea..adf29940a 100755 --- a/scripts/local_p3o_task40/analyze_step0_revalidation.py +++ b/scripts/local_p3o_task40/analyze_step0_revalidation.py @@ -21,6 +21,7 @@ LOGPROB_ABS_TOL = 1e-6 GRAD_REL_L2_TOL = 1e-6 GRAD_COSINE_MIN = 1.0 - 1e-9 +COMMIT_UNDER_TEST = "ef854a3cfb3eb692e083d14e12fbd96a8b1c7af5" TOPOLOGIES = ("dp1", "dp4cp1", "dp2cp1", "dp2cp2") METRIC_KEYS = { "normalized_ess": "train/p3o/normalized_ess", @@ -518,7 +519,7 @@ def main() -> None: route = "Batch 5" if input_identity_failed else ("Batch 6" if failed_stages else None) result = { "batch": 7, - "commit": "e16d325996644f6b205d5b50a189d7e3b35d23da", + "commit": COMMIT_UNDER_TEST, "status": "COMPLETE" if overall_pass else "FAIL_ROUTING", "allow_batch8": overall_pass, "thresholds": { diff --git a/tests/scripts/local_p3o_task40/test_step0_revalidation.py b/tests/scripts/local_p3o_task40/test_step0_revalidation.py index a9a235c35..c097f4ada 100644 --- a/tests/scripts/local_p3o_task40/test_step0_revalidation.py +++ b/tests/scripts/local_p3o_task40/test_step0_revalidation.py @@ -106,3 +106,7 @@ def test_launcher_is_frozen_to_four_bf16_step_scope_cells() -> None: assert "P3O_MICRO_BATCH_SIZE=1" in launcher assert "step0_revalidation.configure" in launcher assert "fp32" not in launcher.lower() + + +def test_retry_verdict_is_pinned_to_the_batch6_loop_commit() -> None: + assert analysis.COMMIT_UNDER_TEST == "ef854a3cfb3eb692e083d14e12fbd96a8b1c7af5" From 347b9ef69b54b761247069f4e486b097c7ea93a1 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Mon, 17 Aug 2026 23:21:34 +0800 Subject: [PATCH 20/30] fix(p3o): preserve partitioned training gradients --- examples/algorithms/p3o/README.md | 27 ++- examples/algorithms/p3o/README_zh.md | 17 +- relax/backends/megatron/cp_utils.py | 53 +++++ relax/backends/megatron/initialize.py | 6 + relax/backends/megatron/loss.py | 24 +- relax/backends/megatron/model.py | 135 +++++++++++- relax/backends/megatron/p3o_step.py | 5 +- .../megatron/test_model_provider_vpp.py | 3 + ...test_p3o_attention_partition_invariance.py | 73 +++++++ .../backends/megatron/test_p3o_initialize.py | 11 + tests/backends/megatron/test_p3o_loss.py | 4 + .../test_p3o_training_partition_invariance.py | 205 ++++++++++++++++++ 12 files changed, 529 insertions(+), 34 deletions(-) create mode 100644 tests/backends/megatron/test_p3o_training_partition_invariance.py diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md index e4206239a..b2b177afe 100644 --- a/examples/algorithms/p3o/README.md +++ b/examples/algorithms/p3o/README.md @@ -87,23 +87,26 @@ reward would define an unvalidated hybrid objective. The reward/verifier name valid for compatible datasets. With `--context-parallel-size > 1`, P3O automatically runs every THD -TransformerLayer 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 QKV, attention, residual, and MLP kernels see the same -token order and shape as CP1. It duplicates full-layer 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. +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. -Batch-invariant kernels cover the final normalization and LM head outside the -wrapped TransformerLayers; P3O fails closed if either mode is absent. +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. 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 diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md index fb681cf7a..395856455 100644 --- a/examples/algorithms/p3o/README_zh.md +++ b/examples/algorithms/p3o/README_zh.md @@ -70,11 +70,11 @@ OPD advantage replacement 或 OPD-only reward 都会形成未经验证的混合 reward/verifier 名称 `P3O_RM_TYPE=mopd` 与 `--use-opd` 训练功能无关，兼容数据集 仍可使用。 -当 `--context-parallel-size > 1` 时，P3O 会自动在每个 THD TransformerLayer 之前 -重建全序列，完成整层计算后再把输出切回各 CP rank。该严格路径对 `micro-batch` 和 -`step` 两种 ESS scope 都生效，使 QKV、attention、残差和 MLP kernel 看到与 CP1 -相同的 token 顺序和形状。它会在每个 CP rank 上重复完整层的计算和 activation， -因此峰值显存和计算量都高于原生 CP；长上下文建议开启整层 activation +当 `--context-parallel-size > 1` 时，P3O 会自动为每个 THD TransformerLayer、最终 +normalization 和 LM head 重建全序列，完成计算后再把输出切回各 CP rank。该严格路径对 +`micro-batch` 和 `step` 两种 ESS scope 都生效，使 forward/backward kernel 看到与 CP1 +相同的 token 顺序和形状。它会在每个 CP rank 上重复全序列计算和 activation，因此 +峰值显存和计算量都高于原生 CP；长上下文建议开启整层 activation recomputation。当前支持的合同是标准 zig-zag THD 且 tensor parallel size 为 1。 如果全序列路径发生 OOM，P3O 会终止并拒绝 CP>1，不会静默回退到数值不等价的 原生 CP kernel order。 @@ -83,9 +83,10 @@ recomputation。当前支持的合同是标准 zig-zag THD 且 tensor parallel s `--batch-invariant-mode`。示例脚本会为配对的 P3O/GRPO run 同时启用两者，并在模型 构建前向 Ray runtime 注入 `NCCL_ALGO=Ring`、 `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` 和 -`CUBLAS_WORKSPACE_CONFIG=:4096:8`。batch-invariant kernel 覆盖 wrapped -TransformerLayer 之外的最终 normalization 与 LM head；缺少任一模式时 P3O 会 -fail closed。 +`CUBLAS_WORKSPACE_CONFIG=:4096:8`。缺少任一模式时 P3O 会 fail closed。P3O 在该模式下 +还会关闭 fused weight-gradient accumulation，因为 +当前随附的 batch-invariant TE GEMM 无法遵守跨多个 micro-batch 的 `main_grad` +累积合同。该绕行使用稳定的 TE/DDP 累积路径，可能降低吞吐或增加瞬时梯度显存。 正式默认值为 G=16、global batch 64、micro-batch 1、rollout batch 4、response length 4096 和 30 个 optimizer step（`--num-rollout 30`）。计划配对 seed 为 42、 diff --git a/relax/backends/megatron/cp_utils.py b/relax/backends/megatron/cp_utils.py index 1e75d316c..a0cb82564 100644 --- a/relax/backends/megatron/cp_utils.py +++ b/relax/backends/megatron/cp_utils.py @@ -764,6 +764,59 @@ def gdn_cp_slice( return torch.cat(pieces, dim=0) +class _P3OReplicatedFullSequenceSlice(torch.autograd.Function): + """Slice replicated full-sequence output with a CP-invariant backward. + + Every strict P3O CP rank runs the same full TransformerLayer, but owns a + different downstream loss shard. Reassemble those local output gradients + before the layer backward so every replica performs the same weight-grad + GEMM. Dividing by CP here is cancelled by the later CP parameter-gradient + SUM and by the full-input gather's reduce-scatter SUM. + """ + + @staticmethod + def forward(ctx, full, cu_seqlens, cp_size, cp_rank, cp_group): + ctx.cu_seqlens = list(cu_seqlens) + ctx.cp_size = int(cp_size) + ctx.cp_rank = int(cp_rank) + ctx.cp_group = cp_group + ctx.full_shape = tuple(full.shape) + return gdn_cp_slice(full, ctx.cu_seqlens, ctx.cp_size, ctx.cp_rank) + + @staticmethod + def backward(ctx, grad_local): + full_grad = grad_local.new_zeros(ctx.full_shape) + local_offset = 0 + for full_start, full_end in zip(ctx.cu_seqlens[:-1], ctx.cu_seqlens[1:], strict=True): + chunk_size = (full_end - full_start) // (2 * ctx.cp_size) + first_start = full_start + ctx.cp_rank * chunk_size + second_start = full_start + (2 * ctx.cp_size - ctx.cp_rank - 1) * chunk_size + full_grad[first_start : first_start + chunk_size] = grad_local[local_offset : local_offset + chunk_size] + local_offset += chunk_size + full_grad[second_start : second_start + chunk_size] = grad_local[local_offset : local_offset + chunk_size] + local_offset += chunk_size + if local_offset != grad_local.shape[0]: + raise RuntimeError( + "P3O replicated CP output gradient does not match packed layout: " + f"consumed={local_offset}, local_tokens={grad_local.shape[0]}" + ) + dist.all_reduce(full_grad, op=dist.ReduceOp.SUM, group=ctx.cp_group) + full_grad.mul_(1.0 / ctx.cp_size) + return full_grad, None, None, None, None + + +def p3o_cp_replicated_slice( + full: torch.Tensor, + cu_seqlens: torch.Tensor | list[int], + cp_size: int, + cp_rank: int, + cp_group: dist.ProcessGroup, +) -> torch.Tensor: + """Slice a strict replicated P3O layer and canonicalize its backward.""" + cu = cu_seqlens if isinstance(cu_seqlens, list) else cu_seqlens.tolist() + return _P3OReplicatedFullSequenceSlice.apply(full, cu, cp_size, cp_rank, cp_group) + + class _AllGatherFullSequence(torch.autograd.Function): """All-gather each CP rank's shard into the full sequence; backward reduce- scatters (sums) the gradient. diff --git a/relax/backends/megatron/initialize.py b/relax/backends/megatron/initialize.py index 175ab61f1..4441cfefb 100644 --- a/relax/backends/megatron/initialize.py +++ b/relax/backends/megatron/initialize.py @@ -31,6 +31,12 @@ def _configure_p3o_partition_invariance(args: Namespace) -> None: raise RuntimeError( "P3O strict partition invariance requires " + " and ".join(missing) + "; refusing a non-equivalent run" ) + # MCore's batch-invariant TE GEMM wrapper does not honor the + # ``accumulate=True`` contract when fused wgrad writes directly to + # ``main_grad``: every micro-batch copies over the previous one. Keep + # the batch-invariant kernels, but use TE/DDP's stable unfused gradient + # accumulation path for the complete optimizer step. + args.gradient_accumulation_fusion = False if batch_invariant_mode: enable_batch_invariant_mode() diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index cdd3f87c8..d67eb3e38 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -968,18 +968,24 @@ def p3o_loss_function( clip_high=getattr(args, "clip_high", 0.2), ) - score_loss = sum_of_sample_mean(terms.score_loss) - adaptive_kl_loss = sum_of_sample_mean(terms.adaptive_kl_loss) + # Keep cancellation-sensitive optimizer-step scalars in FP64. The cast is + # differentiable, so token gradients still flow back in their model dtype; + # only the reduction order is protected from DP/CP partition rounding. + def stable_token_sum(values: torch.Tensor) -> torch.Tensor: + return sum_of_sample_mean(values.to(torch.float64)) + + score_loss = stable_token_sum(terms.score_loss) + adaptive_kl_loss = stable_token_sum(terms.adaptive_kl_loss) # behavior_kl_proxy: sampled-token k3 proxy (1-ESS), not full-vocabulary KL. # Measures concentration of importance ratios via ESS, not distributional shift. - behavior_kl_proxy = sum_of_sample_mean(terms.behavior_kl_proxy) + behavior_kl_proxy = stable_token_sum(terms.behavior_kl_proxy) # cap_fraction: fraction of tokens where adaptive cap binds (ratio > ESS). # Different from PPO's clip_fraction which measures fixed-interval clipping. - cap_fraction = sum_of_sample_mean(terms.cap_hits) - clip_fraction = sum_of_sample_mean(terms.clip_hits) + cap_fraction = stable_token_sum(terms.cap_hits) + clip_fraction = stable_token_sum(terms.clip_hits) entropy = torch.cat(log_probs_and_entropy["entropy"], dim=0) - entropy_loss = sum_of_sample_mean(entropy) + entropy_loss = stable_token_sum(entropy) loss = score_loss + adaptive_kl_loss - args.entropy_coef * entropy_loss @@ -990,7 +996,7 @@ def p3o_loss_function( # behavior KL above and reported under its own key. ref_log_probs = torch.cat(batch["ref_log_probs"], dim=0) reference_kl = compute_approx_kl(log_probs, ref_log_probs, kl_loss_type=args.kl_loss_type) - reference_kl_loss = sum_of_sample_mean(reference_kl) + reference_kl_loss = stable_token_sum(reference_kl) reference_kl_metric = reference_kl_loss.clone().detach() loss = loss + args.kl_loss_coef * reference_kl_loss @@ -999,10 +1005,10 @@ def p3o_loss_function( # Global step scalars are reported as scalar * local_valid_tokens so that the # caller's divide-by-global-token-count recovers the scalar itself. - local_valid_tokens = valid_mask.sum().to(torch.float32) + local_valid_tokens = valid_mask.sum().to(torch.float64) def scaled(value: torch.Tensor) -> torch.Tensor: - return (value.to(torch.float32) * local_valid_tokens).clone().detach() + return (value.to(torch.float64) * local_valid_tokens).clone().detach() reported_loss = { "loss": loss.clone().detach(), diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index ee81eb18d..467f1ad63 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -108,12 +108,13 @@ def _p3o_cp_slice( cu_seqlens: list[int], cp_size: int, cp_rank: int, + cp_group: object, ) -> torch.Tensor: """Return the current rank's THD zig-zag shard from full attention output.""" - from .cp_utils import gdn_cp_slice + from .cp_utils import p3o_cp_replicated_slice - return gdn_cp_slice(full_output, cu_seqlens, cp_size, cp_rank) + return p3o_cp_replicated_slice(full_output, cu_seqlens, cp_size, cp_rank, cp_group) def _canonicalize_p3o_attention_input( @@ -295,7 +296,13 @@ def full_sequence_forward( cu_seqlens_cpu, total_lengths, ) - return _p3o_cp_slice(source_layout_output, cu_seqlens_cpu, cp_size, cp_rank), full_context + return _p3o_cp_slice( + source_layout_output, + cu_seqlens_cpu, + cp_size, + cp_rank, + cp_group, + ), full_context except torch.cuda.OutOfMemoryError as exc: raise RuntimeError(P3O_CP_ATTENTION_OOM_ERROR) from exc finally: @@ -355,6 +362,125 @@ def _find_lm_output_layer(model: torch.nn.Module) -> torch.nn.Module | None: return None +def _find_p3o_post_process_modules( + model: torch.nn.Module, +) -> tuple[torch.nn.Module, torch.nn.Module] | None: + """Find the final normalization and LM head on the last pipeline stage.""" + module = unwrap_model(model) + for _ in range(4): + output_layer = getattr(module, "output_layer", None) + decoder = getattr(module, "decoder", None) + final_layernorm = getattr(decoder, "final_layernorm", None) + if output_layer is not None and not isinstance(output_layer, torch.nn.Identity): + if final_layernorm is None: + raise RuntimeError("P3O strict CP post-processing requires decoder.final_layernorm") + return final_layernorm, output_layer + module = getattr(module, "module", None) or getattr(module, "language_model", None) + if module is None: + return None + return None + + +def _build_p3o_full_sequence_post_process_forward( + module: torch.nn.Module, + packed_seq_params: object, + cp_size: int, + cp_group: object, + cp_rank: int, +) -> Callable: + """Make one token-wise post-process module use the canonical CP1 shape.""" + original_forward = module.forward + cu_seqlens_cpu = getattr(packed_seq_params, "_relax_cu_seqlens_cpu") + total_lengths = getattr(packed_seq_params, "_relax_total_lengths") + pad_multiple = getattr(packed_seq_params, "_relax_attention_pad_multiple") + + @wraps(original_forward) + def full_sequence_forward(hidden_states: torch.Tensor, *args: object, **kwargs: object): + try: + full_hidden_states = _p3o_cp_gather_full(hidden_states, cu_seqlens_cpu, cp_size, cp_group) + canonical_hidden_states, _ = _canonicalize_p3o_attention_input( + full_hidden_states, + cu_seqlens_cpu, + total_lengths, + pad_multiple, + ) + sequence_parallel = getattr(module, "sequence_parallel", None) + if sequence_parallel is not None: + module.sequence_parallel = False + try: + result = original_forward(canonical_hidden_states, *args, **kwargs) + finally: + if sequence_parallel is not None: + module.sequence_parallel = sequence_parallel + + output = result[0] if isinstance(result, tuple) else result + source_layout_output = _restore_p3o_attention_output_layout( + output, + cu_seqlens_cpu, + total_lengths, + ) + local_output = _p3o_cp_slice( + source_layout_output, + cu_seqlens_cpu, + cp_size, + cp_rank, + cp_group, + ) + if isinstance(result, tuple): + return (local_output, *result[1:]) + return local_output + except torch.cuda.OutOfMemoryError as exc: + raise RuntimeError(P3O_CP_ATTENTION_OOM_ERROR) from exc + + return full_sequence_forward + + +@contextmanager +def _p3o_full_sequence_post_process( + args: Namespace, + model: torch.nn.Module, + packed_seq_params: object | None, +) -> Iterator[None]: + """Canonicalize final RMSNorm and LM-head forward/backward under strict + CP.""" + if getattr(args, "advantage_estimator", None) != "p3o" or getattr(args, "context_parallel_size", 1) <= 1: + yield + return + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + raise RuntimeError("P3O strict CP post-processing requires THD packed_seq_params") + modules = _find_p3o_post_process_modules(model) + if modules is None: + yield + return + attention = next((module for module in model.modules() if type(module).__name__ == "SelfAttention"), None) + if attention is None: + raise RuntimeError("P3O strict CP post-processing requires a SelfAttention module") + cp_size, cp_group, cp_rank = _resolve_p3o_attention_cp(attention, packed_seq_params) + if cp_size == 1: + yield + return + + originals: list[tuple[torch.nn.Module, Callable]] = [] + try: + for module in modules: + original_forward = module.forward + module.forward = _build_p3o_full_sequence_post_process_forward( + module, + packed_seq_params, + cp_size, + cp_group, + cp_rank, + ) + originals.append((module, original_forward)) + yield + finally: + for module, original_forward in originals: + try: + del module.forward + except AttributeError: + module.forward = original_forward + + @contextmanager def _bypass_output_layer( model: torch.nn.Module, @@ -1469,7 +1595,8 @@ def forward_step( ) as lm_head_forward: output_tensor = model(**forward_kwargs) else: - output_tensor = model(**forward_kwargs) + with _p3o_full_sequence_post_process(args, model, batch.get("packed_seq_params")): + output_tensor = model(**forward_kwargs) if Envs.ENABLE_ROUTING_REPLAY: os.environ["ROUTING_REPLAY_STAGE"] = old_stage diff --git a/relax/backends/megatron/p3o_step.py b/relax/backends/megatron/p3o_step.py index 54951904b..76439036b 100644 --- a/relax/backends/megatron/p3o_step.py +++ b/relax/backends/megatron/p3o_step.py @@ -201,7 +201,10 @@ def forward_step( inner.pg_collection.cp = mpu.get_dynamic_data_context_parallel_groups(group_size=dynamic_cp_size) try: - output_tensor = model_chunk(**forward_kwargs) + from .model import _p3o_full_sequence_post_process + + with _p3o_full_sequence_post_process(args, model_chunk, batch.get("packed_seq_params")): + output_tensor = model_chunk(**forward_kwargs) finally: if orig_cp_group is not None: inner.pg_collection.cp = orig_cp_group diff --git a/tests/backends/megatron/test_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index e45ec03e8..1b1dcbb8f 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -31,6 +31,7 @@ def __init__(self, model=None): self.vision_dp_when_cp = False self.deterministic_mode = False self.batch_invariant_mode = False + self.gradient_accumulation_fusion = True def finalize(self): self.finalized = True @@ -137,6 +138,7 @@ def _bridge_args(**overrides): "virtual_pipeline_model_parallel_size": 2, "deterministic_mode": True, "batch_invariant_mode": True, + "gradient_accumulation_fusion": False, "context_parallel_size": 1, "expert_model_parallel_size": 1, "expert_tensor_parallel_size": 1, @@ -194,6 +196,7 @@ def test_bridge_provider_receives_partition_invariance_modes(monkeypatch): assert provider.deterministic_mode is True assert provider.batch_invariant_mode is True + assert provider.gradient_accumulation_fusion is False assert provider.finalized diff --git a/tests/backends/megatron/test_p3o_attention_partition_invariance.py b/tests/backends/megatron/test_p3o_attention_partition_invariance.py index 04526c874..9a18cd6a5 100644 --- a/tests/backends/megatron/test_p3o_attention_partition_invariance.py +++ b/tests/backends/megatron/test_p3o_attention_partition_invariance.py @@ -18,6 +18,7 @@ from relax.backends.megatron.model import ( P3O_CP_ATTENTION_OOM_ERROR, _install_p3o_full_sequence_attention, + _p3o_full_sequence_post_process, ) @@ -129,6 +130,34 @@ def forward( return output, context +class _TokenShapePostProcess(torch.nn.Module): + def __init__(self, weight: torch.Tensor, *, returns_tuple: bool = False): + super().__init__() + self.weight = torch.nn.Parameter(weight.clone()) + self.returns_tuple = returns_tuple + self.seen_token_counts: list[int] = [] + + def forward(self, hidden_states: torch.Tensor, *args: object, **kwargs: object): + del args, kwargs + self.seen_token_counts.append(hidden_states.shape[0]) + output = (hidden_states @ self.weight) * (1.0 + hidden_states.shape[0] / 100.0) + return (output, None) if self.returns_tuple else output + + +class _PostProcessRoot(torch.nn.Module): + def __init__( + self, + layer: TransformerLayer, + final_layernorm: torch.nn.Module, + output_layer: torch.nn.Module, + ) -> None: + super().__init__() + self.layer = layer + self.decoder = torch.nn.Module() + self.decoder.final_layernorm = final_layernorm + self.output_layer = output_layer + + def _free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) @@ -237,6 +266,50 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non for parameter, reference_parameter in zip(layer.parameters(), reference_layer.parameters(), strict=True): dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) torch.testing.assert_close(parameter.grad, reference_parameter.grad, atol=1e-10, rtol=1e-10) + + norm_weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) + head_weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) + reference_norm = _TokenShapePostProcess(norm_weight) + reference_head = _TokenShapePostProcess(head_weight, returns_tuple=True) + reference_post_input = reference_output.detach().clone().requires_grad_(True) + reference_post_output, _ = reference_head(reference_norm(reference_post_input)) + (reference_post_output * reference_mask).square().sum().backward() + + norm = _TokenShapePostProcess(norm_weight) + head = _TokenShapePostProcess(head_weight, returns_tuple=True) + post_root = _PostProcessRoot(layer, norm, head) + local_post_input = output.detach().clone().requires_grad_(True) + original_norm_forward = norm.forward + original_head_forward = head.forward + with _p3o_full_sequence_post_process(_p3o_args(scope), post_root, packed_seq_params): + local_post_output, _ = head(norm(local_post_input)) + assert norm.forward == original_norm_forward + assert head.forward == original_head_forward + assert norm.seen_token_counts == [len(canonical_hidden)] + assert head.seen_token_counts == [len(canonical_hidden)] + (local_post_output * local_loss_mask).square().sum().backward() + + expected_full_post_input_grad = torch.zeros_like(full_hidden) + source_offset = 0 + canonical_offset = 0 + for total_length, padded_length in zip(total_lengths, cp_padded_lengths, strict=True): + expected_full_post_input_grad[source_offset : source_offset + total_length] = reference_post_input.grad[ + canonical_offset : canonical_offset + total_length + ] + source_offset += padded_length + canonical_offset += total_length + torch.testing.assert_close( + local_post_input.grad, + gdn_cp_slice(expected_full_post_input_grad, cu_seqlens, world_size, rank), + atol=1e-10, + rtol=1e-10, + ) + for parameter, reference_parameter in ( + (norm.weight, reference_norm.weight), + (head.weight, reference_head.weight), + ): + dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) + torch.testing.assert_close(parameter.grad, reference_parameter.grad, atol=1e-10, rtol=1e-10) finally: dist.destroy_process_group() diff --git a/tests/backends/megatron/test_p3o_initialize.py b/tests/backends/megatron/test_p3o_initialize.py index 3a61c74aa..01ece544e 100644 --- a/tests/backends/megatron/test_p3o_initialize.py +++ b/tests/backends/megatron/test_p3o_initialize.py @@ -39,6 +39,17 @@ def test_p3o_partition_modes_enable_batch_invariant_kernels(monkeypatch: pytest. assert calls == ["enabled"] +def test_p3o_partition_modes_disable_fused_wgrad_accumulation(monkeypatch: pytest.MonkeyPatch) -> None: + """BIK's TE wrapper must not overwrite ``main_grad`` between micro- + batches.""" + monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: None) + args = _args(gradient_accumulation_fusion=True) + + initialize._configure_p3o_partition_invariance(args) + + assert args.gradient_accumulation_fusion is False + + def test_non_p3o_does_not_change_partition_modes(monkeypatch: pytest.MonkeyPatch) -> None: calls = [] monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: calls.append("enabled")) diff --git a/tests/backends/megatron/test_p3o_loss.py b/tests/backends/megatron/test_p3o_loss.py index a28649df7..28c8fae52 100644 --- a/tests/backends/megatron/test_p3o_loss.py +++ b/tests/backends/megatron/test_p3o_loss.py @@ -135,6 +135,10 @@ def capture_context(*args, is_dummy=False): loss, metrics = loss_module.p3o_loss_function(args, batch, torch.zeros(1), torch.sum) assert REQUIRED_P3O_METRICS <= metrics.keys() + assert loss.dtype is torch.float64 + assert metrics["p3o/score_loss"].dtype is torch.float64 + assert metrics["p3o/total_loss"].dtype is torch.float64 + assert metrics["p3o/normalized_ess"].dtype is torch.float64 assert not any(metric.startswith("opd/") for metric in metrics) assert torch.equal(metrics["p3o/reference_kl"], torch.zeros(())) assert not metrics["p3o/reference_kl"].requires_grad diff --git a/tests/backends/megatron/test_p3o_training_partition_invariance.py b/tests/backends/megatron/test_p3o_training_partition_invariance.py new file mode 100644 index 000000000..bdd86592f --- /dev/null +++ b/tests/backends/megatron/test_p3o_training_partition_invariance.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""End-to-end regression for P3O's strict DP/CP training-step contract.""" + +from __future__ import annotations + +import os +import socket +from argparse import Namespace +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from relax.backends.megatron import initialize +from relax.backends.megatron.cp_utils import p3o_cp_replicated_slice + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + return int(sock.getsockname()[1]) + + +class _TinyPolicy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = torch.nn.Embedding(32, 12, dtype=torch.bfloat16) + self.linear_in = torch.nn.Linear(12, 24, bias=False, dtype=torch.bfloat16) + self.linear_out = torch.nn.Linear(24, 12, bias=False, dtype=torch.bfloat16) + self.norm = torch.nn.LayerNorm(12, dtype=torch.bfloat16) + self.lm_head = torch.nn.Linear(12, 32, bias=False, dtype=torch.bfloat16) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + hidden = self.embedding(tokens) + residual = hidden + hidden = self.linear_out(F.silu(self.linear_in(hidden))) + return self.lm_head(self.norm(hidden + residual)).float() + + +def _fixture() -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + samples = [] + for index, length in enumerate((4, 8, 4, 8, 8, 12, 4, 8)): + tokens = (torch.arange(length, dtype=torch.long) * (index + 3) + index + 1) % 31 + behavior = -3.6 + 0.007 * torch.arange(length, dtype=torch.float32) + 0.003 * index + advantages = torch.sin(torch.arange(length, dtype=torch.float32) * 0.7 + index) * (index + 1) + samples.append((tokens, behavior, advantages)) + return samples + + +def _current_log_probs(model: _TinyPolicy, tokens: torch.Tensor) -> torch.Tensor: + logits = model(tokens) + labels = torch.roll(tokens, shifts=-1) + return F.log_softmax(logits, dim=-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) + + +def _cp_token_indices(length: int, cp_size: int, cp_rank: int) -> torch.Tensor: + chunk_size = length // (2 * cp_size) + first = torch.arange(cp_rank * chunk_size, (cp_rank + 1) * chunk_size) + second_rank = 2 * cp_size - cp_rank - 1 + second = torch.arange(second_rank * chunk_size, (second_rank + 1) * chunk_size) + return torch.cat((first, second)) + + +def _context(model: _TinyPolicy, samples: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]) -> float: + ratios = [] + with torch.no_grad(): + for tokens, behavior, _ in samples: + ratios.append(torch.exp(_current_log_probs(model, tokens) - behavior)) + ratio = torch.cat(ratios).double() + return float((ratio.sum().square() / (ratio.square().sum() * ratio.numel())).clamp(max=1.0)) + + +def _worker( + rank: int, + world_size: int, + port: int, + dp_size: int, + cp_size: int, + scope: str, + output_path: str, +) -> None: + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + try: + torch.manual_seed(20260817) + model = _TinyPolicy() + initialize.enable_batch_invariant_mode = lambda: None + args = Namespace( + advantage_estimator="p3o", + batch_invariant_mode=True, + deterministic_mode=True, + gradient_accumulation_fusion=True, + ) + initialize._configure_p3o_partition_invariance(args) + + samples = _fixture() + dp_rank = rank // cp_size + cp_rank = rank % cp_size + cp_group = None + if cp_size > 1: + for group_dp_rank in range(dp_size): + group = dist.new_group(ranks=list(range(group_dp_rank * cp_size, (group_dp_rank + 1) * cp_size))) + if group_dp_rank == dp_rank: + cp_group = group + sample_indices = torch.tensor_split(torch.arange(len(samples)), dp_size)[dp_rank].tolist() + step_cap = _context(model, samples) + main_grads = { + name: torch.zeros_like(parameter, dtype=torch.float32) for name, parameter in model.named_parameters() + } + local_loss = torch.zeros((), dtype=torch.float64) + local_tokens = torch.zeros((), dtype=torch.float64) + + for sample_index in sample_indices: + tokens, behavior, advantages = samples[sample_index] + logits = model(tokens) + labels = torch.roll(tokens, shifts=-1) + if scope == "micro-batch": + full_current = F.log_softmax(logits, dim=-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) + full_ratio = torch.exp(full_current.detach() - behavior) + cap = float( + ( + full_ratio.double().sum().square() / (full_ratio.double().square().sum() * full_ratio.numel()) + ).clamp(max=1.0) + ) + else: + cap = step_cap + if cp_size > 1: + logits = p3o_cp_replicated_slice(logits, [0, tokens.numel()], cp_size, cp_rank, cp_group) + token_indices = _cp_token_indices(tokens.numel(), cp_size, cp_rank) + labels = labels[token_indices] + behavior = behavior[token_indices] + advantages = advantages[token_indices] + current = F.log_softmax(logits, dim=-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) + ratio = torch.exp(current.detach() - behavior) + coefficients = torch.minimum(ratio, ratio.new_tensor(cap)) * advantages + loss = -(coefficients.detach().double() * current.double()).sum() + loss.backward() + local_loss += loss.detach().double() + local_tokens += current.numel() + + for name, parameter in model.named_parameters(): + gradient = parameter.grad.detach().float() + if args.gradient_accumulation_fusion: + # This is the contract of TE's fused ``out=main_grad`` wgrad + # path. MCore BIK ignored ``accumulate=True`` and overwrote. + main_grads[name].copy_(gradient) + else: + main_grads[name].add_(gradient) + parameter.grad = None + + dist.all_reduce(local_loss, op=dist.ReduceOp.SUM) + dist.all_reduce(local_tokens, op=dist.ReduceOp.SUM) + flattened = [] + for gradient in main_grads.values(): + dist.all_reduce(gradient, op=dist.ReduceOp.SUM) + gradient.div_(local_tokens) + flattened.append(gradient.reshape(-1)) + if rank == 0: + torch.save( + { + "gradient": torch.cat(flattened), + "loss": local_loss / local_tokens, + "valid_tokens": local_tokens, + }, + output_path, + ) + finally: + dist.destroy_process_group() + + +def _run_topology(tmp_path: Path, name: str, dp_size: int, cp_size: int, scope: str) -> dict[str, torch.Tensor]: + output_path = tmp_path / f"{scope}_{name}.pt" + world_size = dp_size * cp_size + mp.spawn( + _worker, + args=(world_size, _free_port(), dp_size, cp_size, scope, str(output_path)), + nprocs=world_size, + join=True, + ) + return torch.load(output_path, map_location="cpu", weights_only=True) + + +@pytest.mark.parametrize("scope", ["micro-batch", "step"]) +def test_p3o_full_training_backward_matches_dp1_across_dp_and_cp2(tmp_path: Path, scope: str) -> None: + reference = _run_topology(tmp_path, "dp1", dp_size=1, cp_size=1, scope=scope) + for name, dp_size, cp_size in (("dp2", 2, 1), ("dp4", 4, 1), ("dp2cp2", 2, 2)): + candidate = _run_topology(tmp_path, name, dp_size=dp_size, cp_size=cp_size, scope=scope) + reference_gradient = reference["gradient"].double() + candidate_gradient = candidate["gradient"].double() + relative_l2 = torch.linalg.vector_norm(candidate_gradient - reference_gradient) / torch.linalg.vector_norm( + reference_gradient + ) + cosine = torch.dot(reference_gradient, candidate_gradient) / ( + torch.linalg.vector_norm(reference_gradient) * torch.linalg.vector_norm(candidate_gradient) + ) + + assert relative_l2 <= 1e-6, name + assert cosine >= 1 - 1e-9, name + torch.testing.assert_close(candidate["loss"], reference["loss"], rtol=1e-6, atol=1e-9) + assert candidate["valid_tokens"].item() == reference["valid_tokens"].item() From 6bdf9a20becb3988a2a47571d5430a9b6bdae7e6 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 00:09:51 +0800 Subject: [PATCH 21/30] chore(p3o): pin second retry verdict provenance --- scripts/local_p3o_task40/analyze_step0_revalidation.py | 2 +- tests/scripts/local_p3o_task40/test_step0_revalidation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/local_p3o_task40/analyze_step0_revalidation.py b/scripts/local_p3o_task40/analyze_step0_revalidation.py index adf29940a..d75771bcf 100755 --- a/scripts/local_p3o_task40/analyze_step0_revalidation.py +++ b/scripts/local_p3o_task40/analyze_step0_revalidation.py @@ -21,7 +21,7 @@ LOGPROB_ABS_TOL = 1e-6 GRAD_REL_L2_TOL = 1e-6 GRAD_COSINE_MIN = 1.0 - 1e-9 -COMMIT_UNDER_TEST = "ef854a3cfb3eb692e083d14e12fbd96a8b1c7af5" +COMMIT_UNDER_TEST = "347b9ef69b54b761247069f4e486b097c7ea93a1" TOPOLOGIES = ("dp1", "dp4cp1", "dp2cp1", "dp2cp2") METRIC_KEYS = { "normalized_ess": "train/p3o/normalized_ess", diff --git a/tests/scripts/local_p3o_task40/test_step0_revalidation.py b/tests/scripts/local_p3o_task40/test_step0_revalidation.py index c097f4ada..cd9974405 100644 --- a/tests/scripts/local_p3o_task40/test_step0_revalidation.py +++ b/tests/scripts/local_p3o_task40/test_step0_revalidation.py @@ -109,4 +109,4 @@ def test_launcher_is_frozen_to_four_bf16_step_scope_cells() -> None: def test_retry_verdict_is_pinned_to_the_batch6_loop_commit() -> None: - assert analysis.COMMIT_UNDER_TEST == "ef854a3cfb3eb692e083d14e12fbd96a8b1c7af5" + assert analysis.COMMIT_UNDER_TEST == "347b9ef69b54b761247069f4e486b097c7ea93a1" From 6adf45ad6fd7e464316ba24a92912dc36608fda9 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:48:46 +0800 Subject: [PATCH 22/30] fix(p3o): canonicalize strict DP gradients --- examples/algorithms/p3o/README.md | 14 ++ examples/algorithms/p3o/README_zh.md | 11 ++ relax/backends/megatron/actor.py | 4 + relax/backends/megatron/data.py | 93 +++++++++++++ relax/backends/megatron/initialize.py | 22 ++++ .../local_p3o_task40/step0_revalidation.py | 123 +++++++++++++++++- .../backends/megatron/test_p3o_initialize.py | 43 ++++++ .../test_p3o_training_partition_invariance.py | 41 ++++++ .../test_step0_revalidation.py | 25 ++++ 9 files changed, 374 insertions(+), 2 deletions(-) diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md index b2b177afe..80080eb16 100644 --- a/examples/algorithms/p3o/README.md +++ b/examples/algorithms/p3o/README.md @@ -108,6 +108,20 @@ 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. + 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, diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md index 395856455..e0caf7385 100644 --- a/examples/algorithms/p3o/README_zh.md +++ b/examples/algorithms/p3o/README_zh.md @@ -88,6 +88,17 @@ recomputation。当前支持的合同是标准 zig-zag THD 且 tensor parallel s 当前随附的 batch-invariant TE GEMM 无法遵守跨多个 micro-batch 的 `main_grad` 累积合同。该绕行使用稳定的 TE/DDP 累积路径，可能降低吞吐或增加瞬时梯度显存。 +当 data parallel size 大于 1 时，严格 P3O 还会按全局 DP-rank 顺序在每个 replica +上收集最终训练 batch。DP rank 0 保留真实 loss mask，其余 replica 使用零 loss mask +执行相同的 micro-batch 调度，从而保持全部 DP/CP collective 对齐；既有的 DP×CP 求和 +归约随后传播唯一的 canonical CUDA backward 结果。该路径可消除 rank-local backward +漂移，但会有意重复前后向计算，因此严格模式不提供 DP 训练加速。 +严格路径还会使用 Megatron eager fused-cross-entropy helper：模块导入时已经被 +`torch.compile` 装饰的 wrapper 会被显式解除，因为多个 replica 并发执行同一 +canonical micro-batch 时，其 Triton autotune launcher 不可用。 +当前 strict-DP 合同要求每个 optimizer step 只有一个 rollout mini；如果存在多个 +mini 边界会 fail closed，不会静默重排。 + 正式默认值为 G=16、global batch 64、micro-batch 1、rollout batch 4、response length 4096 和 30 个 optimizer step（`--num-rollout 30`）。计划配对 seed 为 42、 123 和 2026。smoke 使用 G=4、global batch 16、response length 128 和 1 个 diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 6b6057457..aefca3c7d 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -85,6 +85,7 @@ ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY, DataIterator, build_rollout_minibatch_plan, + canonicalize_p3o_strict_dp_rollout, concat_rollout_batches, get_data_iterator, log_perf_data, @@ -914,6 +915,9 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: log_rollout_data(rollout_id, self.args, rollout_data) + if canonicalize_p3o_strict_dp_rollout(self.args, rollout_data): + data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + # Train if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index 508d5075c..8e9ab64ee 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -168,6 +168,99 @@ def concat_rollout_batches(rollout_batches: Sequence[RolloutBatch]) -> RolloutBa return merged +def _cpu_detached_rollout_value(value: Any) -> Any: + """Copy one rollout value into a Gloo-serializable CPU form.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, list): + return [_cpu_detached_rollout_value(item) for item in value] + if isinstance(value, tuple): + return tuple(_cpu_detached_rollout_value(item) for item in value) + if isinstance(value, dict): + return {key: _cpu_detached_rollout_value(item) for key, item in value.items()} + return deepcopy(value) + + +def _merge_p3o_strict_dp_rollout_batches( + rollout_batches: Sequence[RolloutBatch], + *, + is_anchor: bool, +) -> RolloutBatch: + """Build one global strict-P3O batch and keep loss on one DP replica. + + Every DP replica executes the same globally ordered micro-batches so the + schedule and any CP collectives remain aligned. Only DP rank zero keeps the + real loss mask; other replicas contribute an exact zero to the existing + DP×CP gradient reduction. + """ + local_step_counts = [] + for batch in rollout_batches: + counts = batch.get(ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY) + local_sample_count = len(batch.get("total_lengths", [])) + if not isinstance(counts, list) or len(counts) != 1 or counts[0] != local_sample_count: + raise RuntimeError( + "P3O strict DP canonicalization currently requires exactly one rollout mini per optimizer step; " + f"got counts={counts}, local_sample_count={local_sample_count}" + ) + local_step_counts.append(counts[0]) + + merged = concat_rollout_batches(rollout_batches) + global_sample_count = len(merged.get("total_lengths", [])) + if global_sample_count <= 0: + raise ValueError("P3O strict DP canonicalization requires at least one global sample") + merged[ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY] = [sum(local_step_counts)] + if not is_anchor: + loss_masks = merged.get("loss_masks") + if not isinstance(loss_masks, list) or len(loss_masks) != global_sample_count: + raise ValueError("P3O strict DP canonicalization requires one loss mask per sample") + merged["loss_masks"] = [ + torch.zeros_like(mask) if isinstance(mask, torch.Tensor) else [0 for _ in mask] for mask in loss_masks + ] + return merged + + +def canonicalize_p3o_strict_dp_rollout(args: Namespace, rollout_data: RolloutBatch) -> bool: + """Canonicalize strict-P3O training onto DP rank zero. + + Rank-local CUDA backward is deterministic but not invariant across physical + DP rank mappings. Gather the final training batch within each fixed-CP data + parallel group, preserve global sample order, and execute real loss only on + DP rank zero. Existing summed DP×CP gradient and token-count reductions then + propagate exactly one canonical gradient to every optimizer shard. + + Returns ``True`` when the rollout was replaced and its iterator must be + rebuilt, otherwise ``False``. + """ + if getattr(args, "advantage_estimator", None) != "p3o": + return False + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + if dp_size <= 1: + return False + if not getattr(args, "deterministic_mode", False) or not getattr(args, "batch_invariant_mode", False): + raise RuntimeError("P3O strict DP canonicalization requires deterministic and batch-invariant modes") + if getattr(args, "use_routing_replay", False): + raise RuntimeError("P3O strict DP canonicalization does not support routing replay") + + dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False) + dp_gloo_group = mpu.get_data_parallel_group_gloo(with_context_parallel=False) + gathered: list[RolloutBatch | None] = [None for _ in range(dp_size)] + dist.all_gather_object( + gathered, + _cpu_detached_rollout_value(rollout_data), + group=dp_gloo_group, + ) + if any(batch is None for batch in gathered): + raise RuntimeError("P3O strict DP canonicalization did not receive every DP rollout shard") + canonical = _merge_p3o_strict_dp_rollout_batches( + [batch for batch in gathered if batch is not None], + is_anchor=dp_rank == 0, + ) + canonical = move_tensors_to_device(canonical, device_utils.make_current_torch_device()) + rollout_data.clear() + rollout_data.update(canonical) + return True + + PAD_RULES = { # shape like [1, 128, 1036] "input_features": dict( diff --git a/relax/backends/megatron/initialize.py b/relax/backends/megatron/initialize.py index 4441cfefb..ae4bac3cb 100644 --- a/relax/backends/megatron/initialize.py +++ b/relax/backends/megatron/initialize.py @@ -17,6 +17,21 @@ logger = get_logger(__name__) +def _disable_p3o_compiled_fused_cross_entropy() -> None: + """Restore eager fused-CE helpers already decorated at import time.""" + from megatron.core.fusions import fused_cross_entropy + + for name in ( + "calculate_logits_max", + "calculate_predicted_logits", + "calculate_cross_entropy_loss", + "calculate_gradients", + ): + function = getattr(fused_cross_entropy, name) + original = getattr(function, "_torchdynamo_orig_callable", getattr(function, "__wrapped__", function)) + setattr(fused_cross_entropy, name, original) + + def _configure_p3o_partition_invariance(args: Namespace) -> None: """Enable batch-invariant kernels and fail closed for incomplete P3O mode.""" @@ -37,6 +52,13 @@ def _configure_p3o_partition_invariance(args: Namespace) -> None: # the batch-invariant kernels, but use TE/DDP's stable unfused gradient # accumulation path for the complete optimizer step. args.gradient_accumulation_fusion = False + # Canonical DP replicas execute identical shapes concurrently. The + # repository's optional jit_fuser wrapper routes fused CE through a + # TorchInductor autotune path that is not valid for this strict setup; + # use the existing eager fallback while keeping TE/FlashAttention's + # deterministic kernels enabled. + args.disable_jit_fuser = True + _disable_p3o_compiled_fused_cross_entropy() if batch_invariant_mode: enable_batch_invariant_mode() diff --git a/scripts/local_p3o_task40/step0_revalidation.py b/scripts/local_p3o_task40/step0_revalidation.py index 5af1122c2..3220d3948 100755 --- a/scripts/local_p3o_task40/step0_revalidation.py +++ b/scripts/local_p3o_task40/step0_revalidation.py @@ -11,6 +11,7 @@ import hashlib import json +import os import threading from pathlib import Path from typing import Any, Iterator @@ -18,6 +19,13 @@ EXPECTED_FIXTURE_SHA256 = "48538d165386dc94006613d857c022a7ba2e979bdc31bc617374eee2dc3c35b8" FORMAT_VERSION = 1 +DEFAULT_PRE_SYNC_GRADIENT_TARGETS = ( + "decoder.final_layernorm.weight", + "decoder.layers.0.mlp.linear_fc1.layer_norm_weight", + "decoder.layers.0.self_attention.linear_qkv.bias", + "decoder.layers.0.self_attention.linear_qkv.layer_norm_weight", + "decoder.layers.0.self_attention.linear_qkv.weight", +) _STATE: dict[str, Any] = { "args": None, "output_dir": None, @@ -28,6 +36,9 @@ "parameter_capture_error": None, "parameter_capture_thread": None, "gradients_captured": False, + "pre_sync_gradients_captured": False, + "pre_sync_gradient_targets": DEFAULT_PRE_SYNC_GRADIENT_TARGETS, + "strict_dp_anchor": True, "optimizer_ids": set(), } @@ -264,6 +275,82 @@ def _capture_prepared_gradient_shards(model: Any, optimizer: Any) -> None: _STATE["gradients_captured"] = True +def _capture_pre_sync_gradients(model: Any, num_tokens: Any) -> None: + """Capture selected full local gradients before DDP/CP synchronization. + + The normal Batch-7 gradient artifact is intentionally post-finalization. + This diagnostic snapshot keeps a small, named subset of ``main_grad`` + tensors at the finalizer entry so local CUDA backward/accumulation drift + can be separated from the following reduce-scatter/all-reduce and token + normalization. + """ + import torch + + if _STATE["pre_sync_gradients_captured"]: + raise RuntimeError("Task40 Batch 7 attempted to capture pre-sync gradients more than once") + output_dir: Path = _STATE["output_dir"] + rank = int(_STATE["rank"]) + targets = tuple(str(value) for value in _STATE["pre_sync_gradient_targets"]) + gradient_dir = output_dir / "pre_sync_gradients" + vector_dir = gradient_dir / f"rank{rank:05d}" + vector_dir.mkdir(parents=True, exist_ok=False) + + tensors = [] + selected = sorted( + ((name, parameter) for name, parameter in _named_parameters(model) if any(key in name for key in targets)), + key=lambda item: item[0], + ) + if not selected: + raise RuntimeError(f"Task40 Batch 7 pre-sync gradient targets matched no parameters: {targets}") + for index, (name, parameter) in enumerate(selected): + gradient = getattr(parameter, "main_grad", None) + if gradient is None: + gradient = parameter.grad + if gradient is None: + raise RuntimeError(f"Task40 Batch 7 pre-sync gradient is missing for {name}") + cpu, raw = _tensor_bytes(gradient) + values = cpu.to(torch.float64) + file_name = f"tensor_{index:03d}.pt" + torch.save(cpu, vector_dir / file_name) + l2_sq = float(torch.sum(values * values)) + tensors.append( + { + "name": name, + "parameter_shape": list(parameter.shape), + "shape": list(cpu.shape), + "dtype": str(cpu.dtype), + "numel": cpu.numel(), + "finite": bool(torch.isfinite(values).all()), + "l2_sq": l2_sq, + "l2_norm": l2_sq**0.5, + "max_abs": float(values.abs().max()) if values.numel() else 0.0, + "sum": float(values.sum()), + "sha256": hashlib.sha256(raw).hexdigest(), + "file": f"rank{rank:05d}/{file_name}", + } + ) + + local_num_tokens = _cpu(num_tokens) + if isinstance(local_num_tokens, torch.Tensor): + local_num_tokens = int(local_num_tokens.reshape(())) + elif local_num_tokens is not None: + local_num_tokens = int(local_num_tokens) + _write_json( + gradient_dir / f"summary_rank{rank}.json", + { + "format_version": FORMAT_VERSION, + "rank": rank, + "capture_point": "before DDP/CP gradient synchronization and token normalization", + "local_num_tokens": local_num_tokens, + "targets": list(targets), + "tensor_count": len(tensors), + "all_finite": all(record["finite"] for record in tensors), + "tensors": tensors, + }, + ) + _STATE["pre_sync_gradients_captured"] = True + + def configure(args: Any) -> None: """Install the Batch-7 BF16-only oracle observers.""" import torch @@ -294,7 +381,21 @@ def configure(args: Any) -> None: output_dir = Path(args.dump_details).parent / "oracle" output_dir.mkdir(parents=True, exist_ok=True) rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 - _STATE.update({"args": args, "output_dir": output_dir, "rank": rank}) + configured_targets = os.environ.get("P3O_B7_PRE_SYNC_GRADIENT_TARGETS", "").strip() + pre_sync_gradient_targets = ( + tuple(value.strip() for value in configured_targets.split(",") if value.strip()) + if configured_targets + else DEFAULT_PRE_SYNC_GRADIENT_TARGETS + ) + _STATE.update( + { + "args": args, + "output_dir": output_dir, + "rank": rank, + "pre_sync_gradient_targets": pre_sync_gradient_targets, + "strict_dp_anchor": mpu.get_data_parallel_rank(with_context_parallel=False) == 0, + } + ) runtime = { "format_version": FORMAT_VERSION, "rank": rank, @@ -322,6 +423,8 @@ def configure(args: Any) -> None: "joined before optimizer prepare/update" ), "gradient_capture_point": "after DistributedOptimizer.prepare_grads and before clipping/update", + "pre_sync_gradient_capture_point": "before DDP/CP gradient synchronization and token normalization", + "pre_sync_gradient_targets": list(pre_sync_gradient_targets), } _write_json(output_dir / f"runtime_rank{rank}.json", runtime) @@ -331,7 +434,7 @@ def configure(args: Any) -> None: def observed_local_stats(local_args: Any, batch: dict[str, Any], log_probs: list[Any]) -> Any: result = original_local_stats(local_args, batch, log_probs) local_counter = int(_STATE["local_counter"]) - if not batch.get("__is_dummy__", False): + if not batch.get("__is_dummy__", False) and bool(_STATE["strict_dp_anchor"]): current = torch.cat(log_probs, dim=0) behavior = torch.cat(batch["rollout_log_probs"], dim=0) valid_mask = p3o_step.get_cp_local_valid_mask( @@ -412,6 +515,22 @@ def before_train_step( if optimizer_id in _STATE["optimizer_ids"]: raise RuntimeError("Task40 Batch 7 optimizer step observer was already installed") _STATE["optimizer_ids"].add(optimizer_id) + from megatron.core.utils import get_model_config + + config = get_model_config(model[0]) + original_finalize_model_grads = config.finalize_model_grads_func + if original_finalize_model_grads is None: + raise RuntimeError("Task40 Batch 7 requires Megatron's gradient finalizer") + + def observed_finalize_model_grads( + finalizer_model: Any, + num_tokens: Any = None, + **finalizer_kwargs: Any, + ) -> Any: + _capture_pre_sync_gradients(finalizer_model, num_tokens) + return original_finalize_model_grads(finalizer_model, num_tokens, **finalizer_kwargs) + + config.finalize_model_grads_func = observed_finalize_model_grads original_prepare_grads = optimizer.prepare_grads def observed_prepare_grads(*prepare_args: Any, **prepare_kwargs: Any) -> Any: diff --git a/tests/backends/megatron/test_p3o_initialize.py b/tests/backends/megatron/test_p3o_initialize.py index 01ece544e..b8166b5f9 100644 --- a/tests/backends/megatron/test_p3o_initialize.py +++ b/tests/backends/megatron/test_p3o_initialize.py @@ -50,6 +50,49 @@ def test_p3o_partition_modes_disable_fused_wgrad_accumulation(monkeypatch: pytes assert args.gradient_accumulation_fusion is False +def test_p3o_partition_modes_disable_jit_fuser(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: None) + calls = [] + monkeypatch.setattr(initialize, "_disable_p3o_compiled_fused_cross_entropy", lambda: calls.append("disabled")) + args = _args(disable_jit_fuser=False) + + initialize._configure_p3o_partition_invariance(args) + + assert args.disable_jit_fuser is True + assert calls == ["disabled"] + + +def test_p3o_unwraps_fused_cross_entropy_functions(monkeypatch: pytest.MonkeyPatch) -> None: + from megatron.core.fusions import fused_cross_entropy + + def original(*args, **kwargs): + return args, kwargs + + def compiled(*args, **kwargs): + raise AssertionError((args, kwargs)) + + compiled._torchdynamo_orig_callable = original + for name in ( + "calculate_logits_max", + "calculate_predicted_logits", + "calculate_cross_entropy_loss", + "calculate_gradients", + ): + monkeypatch.setattr(fused_cross_entropy, name, compiled) + + initialize._disable_p3o_compiled_fused_cross_entropy() + + assert all( + getattr(fused_cross_entropy, name) is original + for name in ( + "calculate_logits_max", + "calculate_predicted_logits", + "calculate_cross_entropy_loss", + "calculate_gradients", + ) + ) + + def test_non_p3o_does_not_change_partition_modes(monkeypatch: pytest.MonkeyPatch) -> None: calls = [] monkeypatch.setattr(initialize, "enable_batch_invariant_mode", lambda: calls.append("enabled")) diff --git a/tests/backends/megatron/test_p3o_training_partition_invariance.py b/tests/backends/megatron/test_p3o_training_partition_invariance.py index bdd86592f..917ab790c 100644 --- a/tests/backends/megatron/test_p3o_training_partition_invariance.py +++ b/tests/backends/megatron/test_p3o_training_partition_invariance.py @@ -203,3 +203,44 @@ def test_p3o_full_training_backward_matches_dp1_across_dp_and_cp2(tmp_path: Path assert cosine >= 1 - 1e-9, name torch.testing.assert_close(candidate["loss"], reference["loss"], rtol=1e-6, atol=1e-9) assert candidate["valid_tokens"].item() == reference["valid_tokens"].item() + + +def test_p3o_strict_dp_rollout_merge_keeps_one_loss_anchor() -> None: + from relax.backends.megatron.data import _merge_p3o_strict_dp_rollout_batches + + batches = [ + { + "tokens": [torch.tensor([rank + 1, rank + 2])], + "total_lengths": [2], + "response_lengths": [2], + "loss_masks": [torch.ones(2, dtype=torch.int64)], + "advantages": [torch.tensor([float(rank + 1), float(rank + 1)])], + "rollout_mini_local_sample_counts": [1], + "dynamic_global_batch_size": 2, + } + for rank in range(2) + ] + + anchor = _merge_p3o_strict_dp_rollout_batches(batches, is_anchor=True) + replica = _merge_p3o_strict_dp_rollout_batches(batches, is_anchor=False) + + assert anchor["total_lengths"] == [2, 2] + assert anchor["rollout_mini_local_sample_counts"] == [2] + assert [mask.tolist() for mask in anchor["loss_masks"]] == [[1, 1], [1, 1]] + assert [mask.tolist() for mask in replica["loss_masks"]] == [[0, 0], [0, 0]] + assert [value.tolist() for value in replica["advantages"]] == [[1.0, 1.0], [2.0, 2.0]] + + +def test_p3o_strict_dp_rollout_merge_rejects_multiple_rollout_minis() -> None: + from relax.backends.megatron.data import _merge_p3o_strict_dp_rollout_batches + + batch = { + "tokens": [torch.tensor([1]), torch.tensor([2])], + "total_lengths": [1, 1], + "response_lengths": [1, 1], + "loss_masks": [torch.ones(1), torch.ones(1)], + "rollout_mini_local_sample_counts": [1, 1], + } + + with pytest.raises(RuntimeError, match="exactly one rollout mini"): + _merge_p3o_strict_dp_rollout_batches([batch], is_anchor=True) diff --git a/tests/scripts/local_p3o_task40/test_step0_revalidation.py b/tests/scripts/local_p3o_task40/test_step0_revalidation.py index cd9974405..0dc770591 100644 --- a/tests/scripts/local_p3o_task40/test_step0_revalidation.py +++ b/tests/scripts/local_p3o_task40/test_step0_revalidation.py @@ -81,6 +81,31 @@ def test_gradient_comparator_reports_full_vector_rel_l2_and_cosine(tmp_path: Pat assert comparison["full_parameter_cosine"] == 1.0 +def test_pre_sync_gradient_capture_records_unreduced_main_grad(tmp_path: Path, monkeypatch) -> None: + model = torch.nn.Linear(3, 2, bias=True) + model.weight.main_grad = torch.arange(6, dtype=torch.float32).reshape_as(model.weight) + model.bias.main_grad = torch.tensor([7.0, 8.0], dtype=torch.float32) + state = { + **hook._STATE, + "output_dir": tmp_path, + "rank": 3, + "pre_sync_gradients_captured": False, + "pre_sync_gradient_targets": ("weight", "bias"), + } + monkeypatch.setattr(hook, "_STATE", state) + + hook._capture_pre_sync_gradients([model], torch.tensor(11)) + + summary = __import__("json").loads((tmp_path / "pre_sync_gradients" / "summary_rank3.json").read_text()) + assert summary["local_num_tokens"] == 11 + assert summary["capture_point"] == "before DDP/CP gradient synchronization and token normalization" + assert [record["name"] for record in summary["tensors"]] == ["chunk000.bias", "chunk000.weight"] + for record in summary["tensors"]: + saved = torch.load(tmp_path / "pre_sync_gradients" / record["file"], weights_only=True) + expected = model.bias.main_grad if record["name"].endswith("bias") else model.weight.main_grad + assert torch.equal(saved, expected) + + def test_reconstruct_gradient_joins_distributed_optimizer_shards(tmp_path: Path) -> None: records = [] for rank, (start, values) in enumerate(((0, torch.tensor([1.0, 2.0])), (2, torch.tensor([3.0, 4.0])))): From e7f6e4c92f40ed6b6d3a1375f7a868a83c5080fa Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:50:00 +0800 Subject: [PATCH 23/30] fix(p3o): canonicalize strict CP loss gradients --- examples/algorithms/p3o/README.md | 6 + examples/algorithms/p3o/README_zh.md | 5 + relax/backends/megatron/cp_utils.py | 61 +++----- relax/backends/megatron/loss.py | 81 +++++++++- relax/backends/megatron/model.py | 145 ++++++++++++++++-- ...test_p3o_attention_partition_invariance.py | 59 +++++++ tests/backends/megatron/test_p3o_loss.py | 69 +++++++++ 7 files changed, 378 insertions(+), 48 deletions(-) diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md index 80080eb16..39095f8b4 100644 --- a/examples/algorithms/p3o/README.md +++ b/examples/algorithms/p3o/README.md @@ -121,6 +121,12 @@ 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 diff --git a/examples/algorithms/p3o/README_zh.md b/examples/algorithms/p3o/README_zh.md index e0caf7385..6bd26fb91 100644 --- a/examples/algorithms/p3o/README_zh.md +++ b/examples/algorithms/p3o/README_zh.md @@ -98,6 +98,11 @@ recomputation。当前支持的合同是标准 zig-zag THD 且 tensor parallel s canonical micro-batch 时，其 Triton autotune launcher 不可用。 当前 strict-DP 合同要求每个 optimizer step 只有一个 rollout mini；如果存在多个 mini 边界会 fail closed，不会静默重排。 +在每个 CP group 内，严格模式会在 embedding、TransformerLayer、final norm、LM head +和 P3O loss 边界重建 CP1 token 顺序。Forward 值由 CP rank 0 broadcast；backward +把 token-shard 梯度 reduce 到同一 canonical graph，并把其他 replica graph 的梯度 +置零。这样 dense 参数和 embedding 参数都只按 CP1 顺序累加一次，不再平均不同 +rank-local CUDA 或 sparse-embedding 结果。 正式默认值为 G=16、global batch 64、micro-batch 1、rollout batch 4、response length 4096 和 30 个 optimizer step（`--num-rollout 30`）。计划配对 seed 为 42、 diff --git a/relax/backends/megatron/cp_utils.py b/relax/backends/megatron/cp_utils.py index a0cb82564..fb8dc1080 100644 --- a/relax/backends/megatron/cp_utils.py +++ b/relax/backends/megatron/cp_utils.py @@ -764,45 +764,33 @@ def gdn_cp_slice( return torch.cat(pieces, dim=0) -class _P3OReplicatedFullSequenceSlice(torch.autograd.Function): - """Slice replicated full-sequence output with a CP-invariant backward. - - Every strict P3O CP rank runs the same full TransformerLayer, but owns a - different downstream loss shard. Reassemble those local output gradients - before the layer backward so every replica performs the same weight-grad - GEMM. Dividing by CP here is cancelled by the later CP parameter-gradient - SUM and by the full-input gather's reduce-scatter SUM. - """ +class _P3OCanonicalFullSequence(torch.autograd.Function): + """Use CP rank zero's full output and backward graph as the canonical + one.""" @staticmethod - def forward(ctx, full, cu_seqlens, cp_size, cp_rank, cp_group): - ctx.cu_seqlens = list(cu_seqlens) - ctx.cp_size = int(cp_size) - ctx.cp_rank = int(cp_rank) + def forward(ctx, full, cp_group): ctx.cp_group = cp_group - ctx.full_shape = tuple(full.shape) - return gdn_cp_slice(full, ctx.cu_seqlens, ctx.cp_size, ctx.cp_rank) + ctx.cp_rank = dist.get_rank(group=cp_group) + canonical = full.clone() + source = dist.get_global_rank(cp_group, 0) + dist.broadcast(canonical, src=source, group=cp_group) + return canonical @staticmethod - def backward(ctx, grad_local): - full_grad = grad_local.new_zeros(ctx.full_shape) - local_offset = 0 - for full_start, full_end in zip(ctx.cu_seqlens[:-1], ctx.cu_seqlens[1:], strict=True): - chunk_size = (full_end - full_start) // (2 * ctx.cp_size) - first_start = full_start + ctx.cp_rank * chunk_size - second_start = full_start + (2 * ctx.cp_size - ctx.cp_rank - 1) * chunk_size - full_grad[first_start : first_start + chunk_size] = grad_local[local_offset : local_offset + chunk_size] - local_offset += chunk_size - full_grad[second_start : second_start + chunk_size] = grad_local[local_offset : local_offset + chunk_size] - local_offset += chunk_size - if local_offset != grad_local.shape[0]: - raise RuntimeError( - "P3O replicated CP output gradient does not match packed layout: " - f"consumed={local_offset}, local_tokens={grad_local.shape[0]}" - ) - dist.all_reduce(full_grad, op=dist.ReduceOp.SUM, group=ctx.cp_group) - full_grad.mul_(1.0 / ctx.cp_size) - return full_grad, None, None, None, None + def backward(ctx, grad_full): + canonical_grad = grad_full.contiguous() + destination = dist.get_global_rank(ctx.cp_group, 0) + dist.reduce(canonical_grad, dst=destination, op=dist.ReduceOp.SUM, group=ctx.cp_group) + if ctx.cp_rank != 0: + canonical_grad.zero_() + return canonical_grad, None + + +def p3o_cp_canonical_full(full: torch.Tensor, cp_group: dist.ProcessGroup) -> torch.Tensor: + """Broadcast CP0 forward values and reduce the full backward graph to + CP0.""" + return _P3OCanonicalFullSequence.apply(full, cp_group) def p3o_cp_replicated_slice( @@ -812,9 +800,10 @@ def p3o_cp_replicated_slice( cp_rank: int, cp_group: dist.ProcessGroup, ) -> torch.Tensor: - """Slice a strict replicated P3O layer and canonicalize its backward.""" + """Broadcast one canonical full output, then slice this CP rank's shard.""" cu = cu_seqlens if isinstance(cu_seqlens, list) else cu_seqlens.tolist() - return _P3OReplicatedFullSequenceSlice.apply(full, cu, cp_size, cp_rank, cp_group) + canonical = p3o_cp_canonical_full(full, cp_group) + return gdn_cp_slice(canonical, cu, cp_size, cp_rank) class _AllGatherFullSequence(torch.autograd.Function): diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index d67eb3e38..78ac8a59c 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -862,7 +862,7 @@ def get_p3o_context( return finalize_p3o_step_context(stats) -def p3o_loss_function( +def _p3o_loss_function_impl( args: Namespace, batch: RolloutBatch, logits: torch.Tensor, @@ -1035,6 +1035,85 @@ def scaled(value: torch.Tensor) -> torch.Tensor: return loss, reported_loss +def p3o_loss_function( + args: Namespace, + batch: RolloutBatch, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute P3O on one canonical full-response CP loss graph.""" + packed = batch.get("packed_seq_params") + dynamic_cp_size = batch.get("dynamic_cp_size") + dynamic_cp_rank = batch.get("dynamic_cp_rank") + dynamic_cp_group = getattr(packed, "cp_group", None) if dynamic_cp_size is not None else None + cp_size = int(dynamic_cp_size) if dynamic_cp_size is not None else mpu.get_context_parallel_world_size() + full_logits = getattr(packed, "_relax_p3o_canonical_full_logits", None) + if cp_size <= 1 or full_logits is None: + return _p3o_loss_function_impl(args, batch, logits, sum_of_sample_mean) + + total_lengths = [int(value) for value in batch["total_lengths"]] + response_lengths = [int(value) for value in batch["response_lengths"]] + max_seq_lens = batch.get("max_seq_lens") + padded_total_lengths = batch.get("padded_total_lengths") + max_values = max_seq_lens if max_seq_lens is not None else [None] * len(total_lengths) + padded_values = padded_total_lengths if padded_total_lengths is not None else [None] * len(total_lengths) + full_batch = dict(batch) + # Response-valued training fields are CP-local after ``get_batch``. The + # per-sample loss masks intentionally remain full-length and are consumed + # by the CP-aware reducers, so gathering them a second time would violate + # ``all_gather_with_cp``'s local-response input contract. + for field in ("rollout_log_probs", "advantages", "ref_log_probs"): + values = batch.get(field) + if values is None: + continue + full_batch[field] = [ + all_gather_with_cp( + value, + total_length, + response_length, + padded_total_length=padded_total_length, + qkv_format=args.qkv_format, + max_seq_len=max_seq_len, + dynamic_cp_size=dynamic_cp_size, + dynamic_cp_rank=dynamic_cp_rank, + dynamic_cp_group=dynamic_cp_group, + ) + for value, total_length, response_length, max_seq_len, padded_total_length in zip( + values, + total_lengths, + response_lengths, + max_values, + padded_values, + strict=True, + ) + ] + cp_rank = int(dynamic_cp_rank) if dynamic_cp_rank is not None else mpu.get_context_parallel_rank() + if cp_rank != 0: + # Keep micro-batch-scope collectives symmetric without counting the + # replicated full response once per CP rank. CP0 owns the canonical + # loss graph; other ranks retain only the zero-gradient local-logit + # path returned below. + full_batch["loss_masks"] = [torch.zeros_like(mask) for mask in batch["loss_masks"]] + full_batch["dynamic_cp_size"] = 1 + full_batch["dynamic_cp_rank"] = 0 + full_reducer = get_sum_of_sample_mean( + total_lengths, + response_lengths, + full_batch["loss_masks"], + args.calculate_per_token_loss, + args.qkv_format, + max_seq_lens, + padded_total_lengths, + dynamic_cp_size=1, + dynamic_cp_rank=0, + ) + full_loss, full_log = _p3o_loss_function_impl(args, full_batch, full_logits, full_reducer) + if cp_rank == 0: + return full_loss, full_log + zero_loss = 0.0 * logits.sum() + return zero_loss, {key: torch.zeros_like(value) for key, value in full_log.items()} + + def _get_reinforce_plus_plus_mask_safe_reducer( reducer: Callable[[torch.Tensor], torch.Tensor], loss_masks: list[torch.Tensor], diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 467f1ad63..ae5808bee 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -381,12 +381,128 @@ def _find_p3o_post_process_modules( return None +def _find_p3o_embedding(model: torch.nn.Module) -> torch.nn.Module | None: + """Find the token embedding on the first pipeline stage.""" + module = unwrap_model(model) + for _ in range(4): + embedding = getattr(module, "embedding", None) + if embedding is not None and not isinstance(embedding, torch.nn.Identity): + return embedding + module = getattr(module, "module", None) or getattr(module, "language_model", None) + if module is None: + return None + return None + + +def _build_p3o_full_sequence_embedding_forward( + module: torch.nn.Module, + packed_seq_params: object, + cp_size: int, + cp_group: object, + cp_rank: int, +) -> Callable: + """Run the embedding weight-gradient accumulation in canonical CP1 + order.""" + original_forward = module.forward + cu_seqlens_cpu = getattr(packed_seq_params, "_relax_cu_seqlens_cpu") + total_lengths = getattr(packed_seq_params, "_relax_total_lengths") + pad_multiple = getattr(packed_seq_params, "_relax_attention_pad_multiple") + + def canonicalize_batch_first(value: torch.Tensor | None) -> torch.Tensor | None: + if value is None: + return None + if value.ndim < 2 or value.shape[0] != 1: + raise RuntimeError( + f"P3O strict CP embedding requires batch-first singleton inputs, got shape={tuple(value.shape)}" + ) + full_value = _p3o_cp_gather_full(value.transpose(0, 1).contiguous(), cu_seqlens_cpu, cp_size, cp_group) + canonical_value, _ = _canonicalize_p3o_attention_input( + full_value, + cu_seqlens_cpu, + total_lengths, + pad_multiple, + ) + return canonical_value.transpose(0, 1).contiguous() + + @wraps(original_forward) + def full_sequence_forward( + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + tokentype_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + try: + canonical_input_ids = canonicalize_batch_first(input_ids) + if canonical_input_ids is None: + raise RuntimeError("P3O strict CP embedding requires input_ids") + canonical_output = original_forward( + canonical_input_ids, + canonicalize_batch_first(position_ids), + canonicalize_batch_first(tokentype_ids), + ) + from .cp_utils import gdn_cp_slice, p3o_cp_canonical_full + + canonical_output = p3o_cp_canonical_full(canonical_output, cp_group) + source_layout_output = _restore_p3o_attention_output_layout( + canonical_output, + cu_seqlens_cpu, + total_lengths, + ) + return gdn_cp_slice(source_layout_output, cu_seqlens_cpu, cp_size, cp_rank) + except torch.cuda.OutOfMemoryError as exc: + raise RuntimeError(P3O_CP_ATTENTION_OOM_ERROR) from exc + + return full_sequence_forward + + +@contextmanager +def _p3o_full_sequence_embedding( + args: Namespace, + model: torch.nn.Module, + packed_seq_params: object | None, +) -> Iterator[None]: + """Canonicalize token-embedding forward/backward under strict CP.""" + if getattr(args, "advantage_estimator", None) != "p3o" or getattr(args, "context_parallel_size", 1) <= 1: + yield + return + if packed_seq_params is None or getattr(packed_seq_params, "qkv_format", None) != "thd": + raise RuntimeError("P3O strict CP embedding requires THD packed_seq_params") + embedding = _find_p3o_embedding(model) + if embedding is None: + yield + return + attention = next((candidate for candidate in model.modules() if type(candidate).__name__ == "SelfAttention"), None) + if attention is None: + raise RuntimeError("P3O strict CP embedding requires a SelfAttention module") + cp_size, cp_group, cp_rank = _resolve_p3o_attention_cp(attention, packed_seq_params) + if cp_size == 1: + yield + return + + original_forward = embedding.forward + try: + embedding.forward = _build_p3o_full_sequence_embedding_forward( + embedding, + packed_seq_params, + cp_size, + cp_group, + cp_rank, + ) + yield + finally: + try: + del embedding.forward + except AttributeError: + embedding.forward = original_forward + + def _build_p3o_full_sequence_post_process_forward( module: torch.nn.Module, packed_seq_params: object, cp_size: int, cp_group: object, cp_rank: int, + *, + stash_full_logits: bool, ) -> Callable: """Make one token-wise post-process module use the canonical CP1 shape.""" original_forward = module.forward @@ -414,18 +530,23 @@ def full_sequence_forward(hidden_states: torch.Tensor, *args: object, **kwargs: module.sequence_parallel = sequence_parallel output = result[0] if isinstance(result, tuple) else result + from .cp_utils import gdn_cp_slice, p3o_cp_canonical_full + + canonical_output = p3o_cp_canonical_full(output, cp_group) + if stash_full_logits: + # GPTModel's public policy-logit contract is FP32. Stashing + # inside the raw LM head happens before that outer contract is + # applied, so make the canonical full view match the tensor + # consumed by the ordinary CP1 loss path. + packed_seq_params._relax_p3o_canonical_full_logits = ( + canonical_output.transpose(0, 1).float().contiguous() + ) source_layout_output = _restore_p3o_attention_output_layout( - output, + canonical_output, cu_seqlens_cpu, total_lengths, ) - local_output = _p3o_cp_slice( - source_layout_output, - cu_seqlens_cpu, - cp_size, - cp_rank, - cp_group, - ) + local_output = gdn_cp_slice(source_layout_output, cu_seqlens_cpu, cp_size, cp_rank) if isinstance(result, tuple): return (local_output, *result[1:]) return local_output @@ -462,7 +583,7 @@ def _p3o_full_sequence_post_process( originals: list[tuple[torch.nn.Module, Callable]] = [] try: - for module in modules: + for module_index, module in enumerate(modules): original_forward = module.forward module.forward = _build_p3o_full_sequence_post_process_forward( module, @@ -470,6 +591,7 @@ def _p3o_full_sequence_post_process( cp_size, cp_group, cp_rank, + stash_full_logits=module_index == 1, ) originals.append((module, original_forward)) yield @@ -1595,8 +1717,9 @@ def forward_step( ) as lm_head_forward: output_tensor = model(**forward_kwargs) else: - with _p3o_full_sequence_post_process(args, model, batch.get("packed_seq_params")): - output_tensor = model(**forward_kwargs) + with _p3o_full_sequence_embedding(args, model, batch.get("packed_seq_params")): + with _p3o_full_sequence_post_process(args, model, batch.get("packed_seq_params")): + output_tensor = model(**forward_kwargs) if Envs.ENABLE_ROUTING_REPLAY: os.environ["ROUTING_REPLAY_STAGE"] = old_stage diff --git a/tests/backends/megatron/test_p3o_attention_partition_invariance.py b/tests/backends/megatron/test_p3o_attention_partition_invariance.py index 9a18cd6a5..591f42876 100644 --- a/tests/backends/megatron/test_p3o_attention_partition_invariance.py +++ b/tests/backends/megatron/test_p3o_attention_partition_invariance.py @@ -18,6 +18,7 @@ from relax.backends.megatron.model import ( P3O_CP_ATTENTION_OOM_ERROR, _install_p3o_full_sequence_attention, + _p3o_full_sequence_embedding, _p3o_full_sequence_post_process, ) @@ -144,15 +145,35 @@ def forward(self, hidden_states: torch.Tensor, *args: object, **kwargs: object): return (output, None) if self.returns_tuple else output +class _TokenEmbedding(torch.nn.Module): + def __init__(self, weight: torch.Tensor): + super().__init__() + self.word_embeddings = torch.nn.Embedding.from_pretrained(weight.clone(), freeze=False) + self.seen_token_counts: list[int] = [] + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + tokentype_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + del position_ids + assert tokentype_ids is None + self.seen_token_counts.append(input_ids.shape[1]) + return self.word_embeddings(input_ids).transpose(0, 1).contiguous() + + class _PostProcessRoot(torch.nn.Module): def __init__( self, layer: TransformerLayer, final_layernorm: torch.nn.Module, output_layer: torch.nn.Module, + embedding: torch.nn.Module | None = None, ) -> None: super().__init__() self.layer = layer + self.embedding = embedding self.decoder = torch.nn.Module() self.decoder.final_layernorm = final_layernorm self.output_layer = output_layer @@ -267,6 +288,42 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) torch.testing.assert_close(parameter.grad, reference_parameter.grad, atol=1e-10, rtol=1e-10) + embedding_weight = torch.randn(17, hidden_size, dtype=torch.float64) + real_token_ids = [ + torch.tensor([1, 2, 3, 2, 1], dtype=torch.long), + torch.tensor([4, 5, 4, 6, 4], dtype=torch.long), + ] + source_token_ids = torch.cat( + [ + torch.cat([tokens, torch.zeros(padded - len(tokens), dtype=torch.long)]) + for tokens, padded in zip(real_token_ids, cp_padded_lengths, strict=True) + ] + ) + canonical_token_ids = torch.cat([*real_token_ids, torch.zeros(6, dtype=torch.long)]).unsqueeze(0) + reference_embedding = _TokenEmbedding(embedding_weight) + reference_embedding_output = reference_embedding(canonical_token_ids, None) + (reference_embedding_output * reference_mask).square().sum().backward() + + embedding = _TokenEmbedding(embedding_weight) + embedding_root = _PostProcessRoot( + layer, + torch.nn.Identity(), + torch.nn.Identity(), + embedding, + ) + local_token_ids = gdn_cp_slice(source_token_ids, cu_seqlens, world_size, rank).unsqueeze(0) + with _p3o_full_sequence_embedding(_p3o_args(scope), embedding_root, packed_seq_params): + local_embedding_output = embedding(local_token_ids, None) + assert embedding.seen_token_counts == [len(canonical_hidden)] + (local_embedding_output * local_loss_mask).square().sum().backward() + dist.all_reduce(embedding.word_embeddings.weight.grad, op=dist.ReduceOp.SUM) + torch.testing.assert_close( + embedding.word_embeddings.weight.grad, + reference_embedding.word_embeddings.weight.grad, + atol=1e-10, + rtol=1e-10, + ) + norm_weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) head_weight = torch.randn(hidden_size, hidden_size, dtype=torch.float64) reference_norm = _TokenShapePostProcess(norm_weight) @@ -285,6 +342,8 @@ def _cp2_parity_worker(rank: int, world_size: int, port: int, scope: str) -> Non local_post_output, _ = head(norm(local_post_input)) assert norm.forward == original_norm_forward assert head.forward == original_head_forward + assert packed_seq_params._relax_p3o_canonical_full_logits.dtype == torch.float32 + assert packed_seq_params._relax_p3o_canonical_full_logits.shape[:2] == (1, len(canonical_hidden)) assert norm.seen_token_counts == [len(canonical_hidden)] assert head.seen_token_counts == [len(canonical_hidden)] (local_post_output * local_loss_mask).square().sum().backward() diff --git a/tests/backends/megatron/test_p3o_loss.py b/tests/backends/megatron/test_p3o_loss.py index 28c8fae52..c560683ab 100644 --- a/tests/backends/megatron/test_p3o_loss.py +++ b/tests/backends/megatron/test_p3o_loss.py @@ -9,6 +9,7 @@ """ from argparse import Namespace +from types import SimpleNamespace import pytest import torch @@ -193,6 +194,74 @@ def test_p3o_loss_function_normalizes_by_true_valid_tokens(monkeypatch): assert logging_dict["values"][0].item() == 1 +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_p3o_loss_uses_one_canonical_full_cp_graph(monkeypatch, cp_rank): + """Only CP0's full loss survives, while full response masks stay + ungathered.""" + monkeypatch.setattr(loss_module.mpu, "get_context_parallel_world_size", lambda: 2) + monkeypatch.setattr(loss_module.mpu, "get_context_parallel_rank", lambda: cp_rank) + + full_logits = torch.arange(12.0, requires_grad=True).reshape(1, 2, 6) + local_logits = torch.arange(6.0, requires_grad=True).reshape(1, 1, 6) + rollout = torch.tensor([-0.2]) + advantages = torch.tensor([0.5]) + full_mask = torch.tensor([1.0, 0.0]) + gathered_inputs = [] + + def gather(value, *args, **kwargs): + del args, kwargs + gathered_inputs.append(value) + return torch.cat([value, value]) + + reducer_inputs = {} + + def build_reducer(*args, **kwargs): + reducer_inputs["args"] = args + reducer_inputs["kwargs"] = kwargs + return torch.sum + + impl_inputs = {} + + def canonical_impl(args, batch, logits, reducer): + del args + impl_inputs.update(batch=batch, logits=logits, reducer=reducer) + loss = logits.sum() + return loss, {"loss": loss.detach()} + + monkeypatch.setattr(loss_module, "all_gather_with_cp", gather) + monkeypatch.setattr(loss_module, "get_sum_of_sample_mean", build_reducer) + monkeypatch.setattr(loss_module, "_p3o_loss_function_impl", canonical_impl) + batch = { + "packed_seq_params": SimpleNamespace(_relax_p3o_canonical_full_logits=full_logits), + "rollout_log_probs": [rollout], + "advantages": [advantages], + "loss_masks": [full_mask], + "unconcat_tokens": [torch.tensor([1, 2, 3])], + "total_lengths": [3], + "response_lengths": [2], + } + args = Namespace(qkv_format="thd", calculate_per_token_loss=True) + + loss, metrics = loss_module.p3o_loss_function(args, batch, local_logits, torch.sum) + + assert gathered_inputs == [rollout, advantages] + if cp_rank == 0: + assert impl_inputs["batch"]["loss_masks"][0] is full_mask + else: + assert torch.equal(impl_inputs["batch"]["loss_masks"][0], torch.zeros_like(full_mask)) + assert impl_inputs["batch"]["dynamic_cp_size"] == 1 + assert impl_inputs["batch"]["dynamic_cp_rank"] == 0 + assert impl_inputs["logits"] is full_logits + assert reducer_inputs["args"][2][0] is impl_inputs["batch"]["loss_masks"][0] + assert reducer_inputs["kwargs"] == {"dynamic_cp_size": 1, "dynamic_cp_rank": 0} + if cp_rank == 0: + assert torch.equal(loss, full_logits.sum()) + assert torch.equal(metrics["loss"], full_logits.sum().detach()) + else: + assert torch.equal(loss, torch.zeros(())) + assert torch.equal(metrics["loss"], torch.zeros(())) + + def test_policy_loss_dispatch_selects_dedicated_p3o_path(): args = Namespace(advantage_estimator="p3o", use_opd=False) From eed1cb69b9fdb0ed9626929092da12f2e2f38979 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:28:00 +0800 Subject: [PATCH 24/30] chore(p3o): pin successful Batch 7 provenance --- scripts/local_p3o_task40/analyze_step0_revalidation.py | 2 +- tests/scripts/local_p3o_task40/test_step0_revalidation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/local_p3o_task40/analyze_step0_revalidation.py b/scripts/local_p3o_task40/analyze_step0_revalidation.py index d75771bcf..46eeba544 100755 --- a/scripts/local_p3o_task40/analyze_step0_revalidation.py +++ b/scripts/local_p3o_task40/analyze_step0_revalidation.py @@ -21,7 +21,7 @@ LOGPROB_ABS_TOL = 1e-6 GRAD_REL_L2_TOL = 1e-6 GRAD_COSINE_MIN = 1.0 - 1e-9 -COMMIT_UNDER_TEST = "347b9ef69b54b761247069f4e486b097c7ea93a1" +COMMIT_UNDER_TEST = "e7f6e4c92f40ed6b6d3a1375f7a868a83c5080fa" TOPOLOGIES = ("dp1", "dp4cp1", "dp2cp1", "dp2cp2") METRIC_KEYS = { "normalized_ess": "train/p3o/normalized_ess", diff --git a/tests/scripts/local_p3o_task40/test_step0_revalidation.py b/tests/scripts/local_p3o_task40/test_step0_revalidation.py index 0dc770591..0c20d4e3c 100644 --- a/tests/scripts/local_p3o_task40/test_step0_revalidation.py +++ b/tests/scripts/local_p3o_task40/test_step0_revalidation.py @@ -134,4 +134,4 @@ def test_launcher_is_frozen_to_four_bf16_step_scope_cells() -> None: def test_retry_verdict_is_pinned_to_the_batch6_loop_commit() -> None: - assert analysis.COMMIT_UNDER_TEST == "347b9ef69b54b761247069f4e486b097c7ea93a1" + assert analysis.COMMIT_UNDER_TEST == "e7f6e4c92f40ed6b6d3a1375f7a868a83c5080fa" From 6fc8774c2753e39b3f0bc22f5369256be1cf7394 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:24:48 +0800 Subject: [PATCH 25/30] chore(p3o): record replay smoke failure From 4147ab1e1d014ec4fe58a2b549e8842d4516761e Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:31:09 +0800 Subject: [PATCH 26/30] chore(p3o): close clause 2 evidence From 0d0b10cce70094d3b37c8fb2782733690038552f Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:28:53 +0800 Subject: [PATCH 27/30] fix(ci): isolate P3O CPU tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Keep P3O CPU tests importable without Megatron - Scope Megatron, data-loader, and model-registry stubs to CPU-only P3O tests - Preserve synthetic Gloo coverage by resolving bare test roots without Megatron DDP wrappers - Keep fused cross-entropy setup out of the strict partition worker, where it is covered separately --- # 🎨 Style ## Apply required formatter output - Wrap the oracle-vector helper docstring to the configured docformatter width --- .../local_p3o_task40/audit_oracle_vectors.py | 3 +- ...test_p3o_attention_partition_invariance.py | 40 ++++++++++++++----- .../backends/megatron/test_p3o_initialize.py | 25 +++++++++++- .../test_p3o_training_partition_invariance.py | 28 ++++++++++--- 4 files changed, 79 insertions(+), 17 deletions(-) diff --git a/scripts/local_p3o_task40/audit_oracle_vectors.py b/scripts/local_p3o_task40/audit_oracle_vectors.py index f65f95197..f7b96ad5f 100644 --- a/scripts/local_p3o_task40/audit_oracle_vectors.py +++ b/scripts/local_p3o_task40/audit_oracle_vectors.py @@ -61,7 +61,8 @@ def _distribution(values: list[float]) -> dict[str, Any]: def _response_indices( *, total_length: int, response_length: int, cp_size: int, cp_rank: int, max_seq_len: int ) -> list[tuple[int, int]]: - """Return ``(response_index, global_chunk_index)`` in local output order.""" + """Return ``(response_index, global_chunk_index)`` in local output + order.""" if cp_size == 1: return [(index, 0) for index in range(response_length)] diff --git a/tests/backends/megatron/test_p3o_attention_partition_invariance.py b/tests/backends/megatron/test_p3o_attention_partition_invariance.py index 591f42876..833ba7d3a 100644 --- a/tests/backends/megatron/test_p3o_attention_partition_invariance.py +++ b/tests/backends/megatron/test_p3o_attention_partition_invariance.py @@ -7,22 +7,46 @@ import os import socket from argparse import Namespace -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest import torch import torch.distributed as dist import torch.multiprocessing as mp -from relax.backends.megatron.cp_utils import gdn_cp_slice -from relax.backends.megatron.model import ( - P3O_CP_ATTENTION_OOM_ERROR, - _install_p3o_full_sequence_attention, - _p3o_full_sequence_embedding, - _p3o_full_sequence_post_process, +from tests.backends.megatron._megatron_stub import ( + isolated_module_cache, + stubbed_megatron_modules, + temporarily_stub_module, ) +stream_dataloader = ModuleType("relax.utils.data.stream_dataloader") +stream_dataloader.StreamingTQIterator = object +relax_models = ModuleType("relax.models") + +with ( + isolated_module_cache("relax.backends.megatron"), + temporarily_stub_module("relax.utils.data.stream_dataloader", stream_dataloader), + temporarily_stub_module("relax.models", relax_models), + stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), +): + from relax.backends.megatron import model as model_module + from relax.backends.megatron.cp_utils import gdn_cp_slice + + +def _unwrap_test_model(model: torch.nn.Module) -> torch.nn.Module: + """Return the synthetic test root without Megatron DDP unwrapping.""" + return model + + +model_module.unwrap_model = _unwrap_test_model +P3O_CP_ATTENTION_OOM_ERROR = model_module.P3O_CP_ATTENTION_OOM_ERROR +_install_p3o_full_sequence_attention = model_module._install_p3o_full_sequence_attention +_p3o_full_sequence_embedding = model_module._p3o_full_sequence_embedding +_p3o_full_sequence_post_process = model_module._p3o_full_sequence_post_process + + class _LayerRoot(torch.nn.Module): def __init__(self, layer: torch.nn.Module): super().__init__() @@ -403,8 +427,6 @@ def test_p3o_full_sequence_attention_gate_is_p3o_cp_only() -> None: def test_p3o_full_sequence_attention_oom_refuses_native_cp_fallback(monkeypatch: pytest.MonkeyPatch) -> None: - from relax.backends.megatron import model as model_module - cp_group = SimpleNamespace(size=lambda: 2, rank=lambda: 0) layer = TransformerLayer(cp_group, torch.eye(2), torch.eye(2)) leaked_group = SimpleNamespace(size=lambda: 1, rank=lambda: 0) diff --git a/tests/backends/megatron/test_p3o_initialize.py b/tests/backends/megatron/test_p3o_initialize.py index b8166b5f9..84a675d74 100644 --- a/tests/backends/megatron/test_p3o_initialize.py +++ b/tests/backends/megatron/test_p3o_initialize.py @@ -1,11 +1,34 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. from argparse import Namespace +from collections.abc import Iterator +from types import ModuleType import pytest import torch -from relax.backends.megatron import initialize +from tests.backends.megatron._megatron_stub import ( + isolated_module_cache, + stubbed_megatron_modules, + temporarily_stub_module, +) + + +relax_models = ModuleType("relax.models") + +with ( + isolated_module_cache("relax.backends.megatron"), + temporarily_stub_module("relax.models", relax_models), + stubbed_megatron_modules(), +): + from relax.backends.megatron import initialize + + +@pytest.fixture(autouse=True) +def _megatron_import_stub() -> Iterator[None]: + """Keep Megatron's synthetic modules available for dynamic imports.""" + with stubbed_megatron_modules(): + yield def _args(**overrides: object) -> Namespace: diff --git a/tests/backends/megatron/test_p3o_training_partition_invariance.py b/tests/backends/megatron/test_p3o_training_partition_invariance.py index 917ab790c..a996bdfeb 100644 --- a/tests/backends/megatron/test_p3o_training_partition_invariance.py +++ b/tests/backends/megatron/test_p3o_training_partition_invariance.py @@ -8,6 +8,7 @@ import socket from argparse import Namespace from pathlib import Path +from types import ModuleType import pytest import torch @@ -15,8 +16,26 @@ import torch.multiprocessing as mp import torch.nn.functional as F -from relax.backends.megatron import initialize -from relax.backends.megatron.cp_utils import p3o_cp_replicated_slice +from tests.backends.megatron._megatron_stub import ( + isolated_module_cache, + stubbed_megatron_modules, + temporarily_stub_module, +) + + +stream_dataloader = ModuleType("relax.utils.data.stream_dataloader") +stream_dataloader.StreamingTQIterator = object +relax_models = ModuleType("relax.models") + +with ( + isolated_module_cache("relax.backends.megatron"), + temporarily_stub_module("relax.utils.data.stream_dataloader", stream_dataloader), + temporarily_stub_module("relax.models", relax_models), + stubbed_megatron_modules(("megatron", "ray", "tensordict", "pybase64")), +): + from relax.backends.megatron import initialize + from relax.backends.megatron.cp_utils import p3o_cp_replicated_slice + from relax.backends.megatron.data import _merge_p3o_strict_dp_rollout_batches def _free_port() -> int: @@ -90,6 +109,7 @@ def _worker( torch.manual_seed(20260817) model = _TinyPolicy() initialize.enable_batch_invariant_mode = lambda: None + initialize._disable_p3o_compiled_fused_cross_entropy = lambda: None args = Namespace( advantage_estimator="p3o", batch_invariant_mode=True, @@ -206,8 +226,6 @@ def test_p3o_full_training_backward_matches_dp1_across_dp_and_cp2(tmp_path: Path def test_p3o_strict_dp_rollout_merge_keeps_one_loss_anchor() -> None: - from relax.backends.megatron.data import _merge_p3o_strict_dp_rollout_batches - batches = [ { "tokens": [torch.tensor([rank + 1, rank + 2])], @@ -232,8 +250,6 @@ def test_p3o_strict_dp_rollout_merge_keeps_one_loss_anchor() -> None: def test_p3o_strict_dp_rollout_merge_rejects_multiple_rollout_minis() -> None: - from relax.backends.megatron.data import _merge_p3o_strict_dp_rollout_batches - batch = { "tokens": [torch.tensor([1]), torch.tensor([2])], "total_lengths": [1, 1], From e4acbb9f82cab030fa2902fb2c33f1efdbcd4b83 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:59:00 +0800 Subject: [PATCH 28/30] fix(ci): track Step-0 oracle analyzer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Ship the Step-0 analysis dependency - Add the oracle analyzer imported by the checked-in revalidation script and its tests - Preserve the existing artifact parsing and token-key semantics used by the audit CLI - Keep the CPU CI collection path self-contained outside local ignored files --- .../local_p3o_task40/analyze_step0_oracle.py | 236 ++++++++++++++++++ 1 file changed, 236 insertions(+) create mode 100755 scripts/local_p3o_task40/analyze_step0_oracle.py diff --git a/scripts/local_p3o_task40/analyze_step0_oracle.py b/scripts/local_p3o_task40/analyze_step0_oracle.py new file mode 100755 index 000000000..7a8f4b016 --- /dev/null +++ b/scripts/local_p3o_task40/analyze_step0_oracle.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 + +"""Analyze fixed-input Step-0 topology runs and emit an auditable verdict.""" + +import argparse +import ast +import hashlib +import json +import math +import re +from pathlib import Path + +import torch + + +ANSI = re.compile(r"\x1b\[[0-9;]*m") +STEP_METRICS = re.compile(r"step 0: (\{.*\})") + + +def _token_key(tokens: torch.Tensor) -> str: + return hashlib.sha256(tokens.to(torch.int64).numpy().tobytes()).hexdigest() + + +def _relative_error(left: float, right: float) -> float: + return abs(left - right) / max(abs(left), abs(right), 1e-30) + + +def _metrics(run_dir: Path) -> dict[str, float]: + text = ANSI.sub("", (run_dir / "stdout_stderr.log").read_text(errors="replace")) + matches = STEP_METRICS.findall(text) + if len(matches) != 1: + raise ValueError(f"expected one Step-0 metric record in {run_dir}, found {len(matches)}") + return ast.literal_eval(matches[0]) + + +def _load_run(run_dir: Path) -> dict: + runtimes = [json.loads(path.read_text()) for path in sorted((run_dir / "oracle").glob("runtime_rank*.json"))] + vector_paths = sorted((run_dir / "oracle").glob("vectors_rank*_micro*.pt")) + global_paths = sorted((run_dir / "oracle").glob("global_stats_rank*_sync0.pt")) + if not runtimes or not vector_paths or not global_paths: + raise ValueError(f"missing oracle artifacts in {run_dir}") + + canonical = {} + seen_tokens = set() + position_ids_valid = True + local_s1 = 0.0 + local_s2 = 0.0 + local_n = 0.0 + runtime_by_rank = {record["rank"]: record for record in runtimes} + for path in vector_paths: + artifact = torch.load(path, map_location="cpu", weights_only=False) + runtime = runtime_by_rank[artifact["rank"]] + cp_rank = runtime["cp_rank"] + cp_size = runtime["cp_world_size"] + qkv_format = runtime.get("qkv_format", "thd") + current = artifact["current_log_probs"].to(torch.float64) + behavior = artifact["rollout_log_probs"].to(torch.float64) + valid = artifact["valid_mask"].bool() + offset = 0 + for tokens, positions, response_length, total_length in zip( + artifact["token_ids"], + artifact["position_ids"], + artifact["response_lengths"], + artifact["total_lengths"], + strict=True, + ): + token_key = _token_key(tokens) + seen_tokens.add(token_key) + position_ids_valid &= torch.equal(positions, torch.arange(len(tokens), dtype=torch.int64)) + if cp_size == 1: + response_indices = list(range(response_length)) + else: + prompt_length = total_length - response_length + if qkv_format == "bshd": + max_seq_lens = artifact.get("max_seq_lens") + if not max_seq_lens: + raise ValueError(f"missing max_seq_lens for BSHD artifact {path}") + max_seq_len = int(max_seq_lens[0]) + chunk_size = (max_seq_len + 2 * cp_size - 1) // (2 * cp_size) + else: + chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size) + token_ranges = ( + ( + max(cp_rank * chunk_size, prompt_length - 1) + 1, + min((cp_rank + 1) * chunk_size, total_length - 1) + 1, + ), + ( + max((2 * cp_size - cp_rank - 1) * chunk_size, prompt_length - 1) + 1, + min((2 * cp_size - cp_rank) * chunk_size, total_length - 1) + 1, + ), + ) + response_indices = [ + token_index - prompt_length + for start, stop in token_ranges + for token_index in range(start, max(start, stop)) + if prompt_length <= token_index < total_length + ] + for local_index, response_index in enumerate(response_indices): + flat_index = offset + local_index + if bool(valid[flat_index]): + key = (token_key, response_index) + if key in canonical: + raise ValueError(f"duplicate valid token key {key} in {run_dir}") + canonical[key] = (float(current[flat_index]), float(behavior[flat_index])) + offset += len(response_indices) + if offset != current.numel() or current.shape != behavior.shape or current.shape != valid.shape: + raise ValueError(f"unaligned token vectors in {path}") + local_s1 += float(artifact["local_s1"]) + local_s2 += float(artifact["local_s2"]) + local_n += float(artifact["local_n"]) + + current = torch.tensor([canonical[key][0] for key in sorted(canonical)], dtype=torch.float64) + behavior = torch.tensor([canonical[key][1] for key in sorted(canonical)], dtype=torch.float64) + # Match production exactly: subtraction is performed in float32 before + # ESS moments are promoted to float64. + log_ratio = (current.to(torch.float32) - behavior.to(torch.float32)).to(torch.float64) + ratio = torch.exp(log_ratio) + recomputed = { + "s1": float(ratio.sum()), + "s2": float(ratio.square().sum()), + "n": int(ratio.numel()), + "kl_sum": float((log_ratio + torch.exp(torch.clamp(-log_ratio, min=-10.0, max=10.0)) - 1.0).sum()), + } + recomputed["normalizer"] = recomputed["n"] + recomputed["ess"] = recomputed["s1"] ** 2 / (recomputed["n"] * (recomputed["s2"] + 1e-8)) + + global_records = [torch.load(path, map_location="cpu", weights_only=False) for path in global_paths] + global_stats = {key: float(global_records[0][key]) for key in ("s1", "s2", "n")} + global_rank_agreement = all( + all(float(record[key]) == global_stats[key] for key in global_stats) for record in global_records + ) + metrics = _metrics(run_dir) + return { + "run_dir": str(run_dir.resolve()), + "exit_code": int((run_dir / "exit_code.txt").read_text().strip()), + "runtime": runtimes, + "vector_artifact_count": len(vector_paths), + "token_stream_count": len(seen_tokens), + "position_ids_valid": position_ids_valid, + "canonical": canonical, + "current": current, + "behavior": behavior, + "local_stats_sum": {"s1": local_s1, "s2": local_s2, "n": local_n}, + "global_stats": global_stats, + "global_rank_agreement": global_rank_agreement, + "recomputed": recomputed, + "metrics": metrics, + } + + +def _public(run: dict) -> dict: + return {key: value for key, value in run.items() if key not in {"canonical", "current", "behavior"}} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dp4-bf16", type=Path, required=True) + parser.add_argument("--dp2cp2-bf16", type=Path, required=True) + parser.add_argument("--dp4-fp32", type=Path, required=True) + parser.add_argument("--dp2cp2-fp32", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + runs = { + "dp4_bf16": _load_run(args.dp4_bf16), + "dp2cp2_bf16": _load_run(args.dp2cp2_bf16), + "dp4_fp32": _load_run(args.dp4_fp32), + "dp2cp2_fp32": _load_run(args.dp2cp2_fp32), + } + comparisons = {} + for precision in ("bf16", "fp32"): + left = runs[f"dp4_{precision}"] + right = runs[f"dp2cp2_{precision}"] + key_match = set(left["canonical"]) == set(right["canonical"]) + behavior_max_abs = float(torch.max(torch.abs(left["behavior"] - right["behavior"]))) if key_match else math.inf + current_max_abs = float(torch.max(torch.abs(left["current"] - right["current"]))) if key_match else math.inf + metric_keys = { + "ess": "train/p3o/normalized_ess", + "cap": "train/p3o/adaptive_cap", + "loss": "train/loss", + "gradient_l2_norm": "train/grad_norm", + } + comparisons[precision] = { + "valid_token_keys_identical": key_match, + "behavior_log_probs_max_abs": behavior_max_abs, + "current_log_probs_max_abs": current_max_abs, + "s1_relative_error": _relative_error(left["recomputed"]["s1"], right["recomputed"]["s1"]), + "s2_relative_error": _relative_error(left["recomputed"]["s2"], right["recomputed"]["s2"]), + "kl_sum_relative_error": _relative_error(left["recomputed"]["kl_sum"], right["recomputed"]["kl_sum"]), + **{ + f"{name}_relative_error": _relative_error(left["metrics"][metric], right["metrics"][metric]) + for name, metric in metric_keys.items() + }, + } + + required_artifacts_pass = all( + run["exit_code"] == 0 + and run["position_ids_valid"] + and run["global_rank_agreement"] + and run["recomputed"]["n"] == 13635 + and abs(run["global_stats"]["s1"] - run["recomputed"]["s1"]) < 1e-9 + and abs(run["global_stats"]["s2"] - run["recomputed"]["s2"]) < 1e-9 + for run in runs.values() + ) + fp32 = comparisons["fp32"] + fp32_alignment_pass = ( + fp32["valid_token_keys_identical"] + and fp32["behavior_log_probs_max_abs"] == 0.0 + and fp32["current_log_probs_max_abs"] <= 1e-7 + and fp32["ess_relative_error"] <= 1e-7 + and fp32["loss_relative_error"] <= 1e-7 + and fp32["gradient_l2_norm_relative_error"] <= 1e-7 + ) + verdict = { + "criterion": "FP32 DP4CP1 vs DP2CP2 aligns to approximately 1e-7; BF16-only differences are numerical noise", + "required_artifacts_pass": required_artifacts_pass, + "fp32_alignment_pass": fp32_alignment_pass, + "classification": ( + "bf16_numerical_noise" + if required_artifacts_pass and fp32_alignment_pass + else "cp_forward_mismatch_remains" + ), + "minimum_p0_passed": required_artifacts_pass and fp32_alignment_pass, + } + output = { + "runs": {name: _public(run) for name, run in runs.items()}, + "comparisons": comparisons, + "verdict": verdict, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n") + print(json.dumps(verdict, sort_keys=True)) + + +if __name__ == "__main__": + main() From 59ec10d30693adcde6571c715d6d16a1dc78b52f Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:07:24 +0800 Subject: [PATCH 29/30] fix(ci): set P3O loss CP stub MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Model the CPU test topology explicitly - Set the context-parallel size to one for the P3O loss schema test - Avoid treating the generic Megatron import stub as a production runtime value --- tests/backends/megatron/test_p3o_loss.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/backends/megatron/test_p3o_loss.py b/tests/backends/megatron/test_p3o_loss.py index c560683ab..b34301143 100644 --- a/tests/backends/megatron/test_p3o_loss.py +++ b/tests/backends/megatron/test_p3o_loss.py @@ -83,6 +83,7 @@ def test_get_p3o_context_rejects_unknown_scope(): def test_p3o_loss_reports_complete_schema_without_reference_kl(monkeypatch): + monkeypatch.setattr(loss_module.mpu, "get_context_parallel_world_size", lambda: 1) step_context = P3OStepContext( normalized_ess=torch.tensor(0.75, dtype=torch.float64), adaptive_cap=torch.tensor(0.75, dtype=torch.float64), From ed5669affd5ed1af811f9e933fd04520641017c9 Mon Sep 17 00:00:00 2001 From: DreamEnding <63937131+DreamEnding@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:17:17 +0800 Subject: [PATCH 30/30] fix(ci): stub lazy P3O model import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Keep P3O step kwargs tests CPU-safe - Provide the no-op post-process context required by the delayed model import - Restore the temporary model module after each test to avoid cache leakage --- tests/backends/megatron/test_p3o_step.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/backends/megatron/test_p3o_step.py b/tests/backends/megatron/test_p3o_step.py index dab2b6b02..581d14ce6 100644 --- a/tests/backends/megatron/test_p3o_step.py +++ b/tests/backends/megatron/test_p3o_step.py @@ -5,13 +5,15 @@ from __future__ import annotations import sys -from types import SimpleNamespace +from collections.abc import Iterator +from contextlib import nullcontext +from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock import pytest import torch -from tests.backends.megatron._megatron_stub import stubbed_megatron_modules +from tests.backends.megatron._megatron_stub import stubbed_megatron_modules, temporarily_stub_module with stubbed_megatron_modules(("megatron", "ray", "tensordict")): @@ -34,6 +36,14 @@ def _stub_cp_world_size(monkeypatch): monkeypatch.setattr(p3o_step.mpu, "is_pipeline_last_stage", lambda ignore_virtual=False: True) +@pytest.fixture(autouse=True) +def _stub_p3o_model_import() -> Iterator[None]: + model_module = ModuleType("relax.backends.megatron.model") + model_module._p3o_full_sequence_post_process = lambda *_args, **_kwargs: nullcontext() + with temporarily_stub_module("relax.backends.megatron.model", model_module): + yield + + def _stats(values: tuple[float, float, float]) -> P3OSufficientStats: vector = torch.tensor(values, dtype=torch.float64) return P3OSufficientStats.from_vector(vector)