diff --git a/examples/algorithms/p3o/README.md b/examples/algorithms/p3o/README.md new file mode 100644 index 000000000..63a1d7f2f --- /dev/null +++ b/examples/algorithms/p3o/README.md @@ -0,0 +1,156 @@ +# 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. + +## 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` remains available to the pure full-vocabulary +verification helper, but production argument validation rejects it 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/__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..5739a6998 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -702,6 +702,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..a144fde7f 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 @@ -48,6 +49,7 @@ from .data import DataIterator, 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. @@ -1126,15 +1145,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 +1168,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 +1489,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..bd2ce54ee --- /dev/null +++ b/relax/backends/megatron/p3o_step.py @@ -0,0 +1,332 @@ +# 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, 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), + ) + # The forward inputs must be selected exactly as the training pass in + # model.py::train_one_step does, or the two passes read different token + # layouts and the frozen cap would be computed from logits the gradient + # pass never sees. The VL bridge (Qwen3VLModel.forward) does its own + # CP+SP splitting, so it takes unsplit tokens and no caller-side + # packed_seq_params. + # + # NOTE: This logic is intentionally duplicated from model.py::train_one_step + # rather than extracted to a shared helper. The duplication ensures that + # any future changes to model.py's forward input preparation are immediately + # visible as a diff here, preventing silent drift between the stats pass + # and the training pass. If this block and model.py diverge, the ESS cap + # is computed from different logits than the gradient, breaking P3O. + mm_inputs = batch.get("multimodal_train_inputs") + mm_kwargs = mm_inputs if getattr(args, "is_vl_model", False) and mm_inputs else {} + needs_unsplit = ( + getattr(args, "is_vl_model", False) + or batch.get("multimodal_train_inputs") is not None + or getattr(args, "uses_unsplit_forward", False) + ) + + if needs_unsplit and "unsplit_tokens" in batch: + forward_input_ids = batch["unsplit_tokens"] + forward_packed_seq_params = None + else: + forward_input_ids = batch["tokens"] + forward_packed_seq_params = batch["packed_seq_params"] + + # thd bridge+CP: the bridge needs the per-sample attention mask and the + # matching thd packed_seq_params; loss_mask is None there because + # labels=None means the model runs no internal loss. + if needs_unsplit and "vlm_packed_seq_params" in batch: + forward_attention_mask = batch["unsplit_attention_mask"] + forward_packed_seq_params = batch["vlm_packed_seq_params"] + forward_loss_mask = None + else: + forward_attention_mask = None + forward_loss_mask = batch["full_loss_masks"] + + # 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( + input_ids=forward_input_ids, + position_ids=None, + attention_mask=forward_attention_mask, + labels=None, + packed_seq_params=forward_packed_seq_params, + loss_mask=forward_loss_mask, + **mm_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/backends/sglang/deterministic_sampler_patch.py b/relax/backends/sglang/deterministic_sampler_patch.py new file mode 100644 index 000000000..1226053ca --- /dev/null +++ b/relax/backends/sglang/deterministic_sampler_patch.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Backport SGLang's deterministic-sampler uint32 endpoint fix. + +SGLang 0.5.12.post1 maps a 32-bit hash to ``[0, 1]`` by dividing by +``uint32.max``. A hash equal to ``0xffffffff`` therefore produces exactly ``x +== 1`` and Gumbel noise ``-log(-log(x)) == +inf``. That token then wins the +argmax regardless of its model probability. Upstream clamps ``log(x)`` away +from zero by one hash bucket; this module applies the same correction in the +scheduler subprocess for the affected local runtime. +""" + +from __future__ import annotations + +from collections.abc import Callable +from importlib.metadata import version +from inspect import signature +from typing import Any + +import torch +from packaging.version import Version + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +_AFFECTED_SGLANG_VERSIONS = frozenset({"0.5.12.post1"}) +_PATCH_MARKER = "_relax_uint32_endpoint_fix" + + +def _installed_sglang_version() -> str: + return Version(version("sglang")).public + + +def _uniform_hash_to_gumbel_(values: torch.Tensor) -> torch.Tensor: + """Transform uniform hash fractions in place without infinite endpoints.""" + values.log_().clamp_(min=torch.finfo(values.dtype).min, max=-(2.0**-32)).neg_() + values.log_().neg_() + return values + + +def _build_safe_multinomial_with_seed( + murmur_hash32: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], + *, + compile_function: Callable[..., Any] = torch.compile, +) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]: + """Build the upstream-equivalent deterministic multinomial function.""" + + def _safe_multinomial_with_seed( + logprobs: torch.Tensor, seed: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + _, vocabulary_size = logprobs.shape + seed = seed.to(torch.uint64) + column_indices = torch.arange(vocabulary_size, device=logprobs.device) + hashed = murmur_hash32(seed, positions, column_indices) + gumbel = hashed.to(torch.float64) / torch.iinfo(torch.uint32).max + _uniform_hash_to_gumbel_(gumbel) + gumbel.add_(logprobs.to(torch.float64)) + return torch.argmax(gumbel, dim=1, keepdim=True) + + return compile_function(dynamic=True)(_safe_multinomial_with_seed) + + +def apply_deterministic_sampler_endpoint_patch() -> bool: + """Patch the affected SGLang sampler before scheduler model initialization. + + Returns ``True`` only when this call installs the backport. Unaffected + versions and already-patched scheduler processes are left unchanged. + """ + installed_version = _installed_sglang_version() + if installed_version not in _AFFECTED_SGLANG_VERSIONS: + logger.info( + "SGLang deterministic sampler endpoint backport not required for version %s", + installed_version, + ) + return False + + from sglang.srt.layers import sampler + from sglang.srt.layers.utils.hash import murmur_hash32 + + current = sampler.multinomial_with_seed + if getattr(current, _PATCH_MARKER, False): + return False + parameter_names = tuple(signature(current).parameters) + if parameter_names != ("logprobs", "seed", "positions"): + raise RuntimeError( + "Affected SGLang multinomial_with_seed signature changed: " + f"expected=('logprobs', 'seed', 'positions'): actual={parameter_names}" + ) + + replacement = _build_safe_multinomial_with_seed(murmur_hash32) + setattr(replacement, _PATCH_MARKER, True) + sampler.multinomial_with_seed = replacement + logger.warning( + "Applied SGLang %s deterministic sampler uint32 endpoint backport", + installed_version, + ) + return True diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 15ca1d172..b167f73a3 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -162,16 +162,22 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int: def _patched_run_scheduler_process(*args, **kwargs): - """Scheduler-subprocess entry used for the routing-replay path. + """Scheduler-subprocess entry for Relax's SGLang runtime patches. - This wrapper is only installed when ``--optimize-routing-replay`` is - enabled (see ``_launch_server_with_patches``), so the routing-replay async - D→H patch is applied **unconditionally** here, preserving the original - behavior. + The deterministic sampler endpoint backport is version-gated internally. + The routing-replay async D→H patch remains gated by its existing runtime + environment flag. """ - from relax.backends.sglang.routing_replay_patch import apply_patch + from relax.backends.sglang.deterministic_sampler_patch import ( + apply_deterministic_sampler_endpoint_patch, + ) + + apply_deterministic_sampler_endpoint_patch() + + if Envs.RELAX_OPTIMIZE_ROUTING_REPLAY: + from relax.backends.sglang.routing_replay_patch import apply_patch - apply_patch() + apply_patch() from sglang.srt.managers.scheduler import run_scheduler_process @@ -184,9 +190,8 @@ def _launch_server_with_patches(server_args: ServerArgs): - main process: OPD pre-expanded multimodal patch (``RELAX_OPD_PREEXPANDED_PATCH=1``). - - scheduler subprocess: routing-replay (``RELAX_OPTIMIZE_ROUTING_REPLAY=1``) - installs ``_patched_run_scheduler_process``, which applies the - routing-replay patch unconditionally. + - scheduler subprocess: version-gated deterministic-sampler endpoint fix; + routing replay remains gated by ``RELAX_OPTIMIZE_ROUTING_REPLAY=1``. """ from sglang.srt.entrypoints.http_server import launch_server @@ -195,10 +200,7 @@ def _launch_server_with_patches(server_args: ServerArgs): apply_opd_preexpanded_patch() - if Envs.RELAX_OPTIMIZE_ROUTING_REPLAY: - launch_server(server_args, run_scheduler_process_func=_patched_run_scheduler_process) - else: - launch_server(server_args) + launch_server(server_args, run_scheduler_process_func=_patched_run_scheduler_process) def _resolve_external_model_arch(package_name): @@ -234,9 +236,9 @@ def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: multiprocessing.set_start_method("spawn", force=True) server_args.host = server_args.host.strip("[]") - # Each SGLang patch is controlled by its own env flag and applied - # independently (see ``_launch_server_with_patches`` and - # ``_patched_run_scheduler_process``); any combination is valid: + # Runtime patches are applied independently in the scheduler subprocess + # (see ``_launch_server_with_patches`` and ``_patched_run_scheduler_process``): + # - deterministic sampler endpoint fix: version-gated backport # - RELAX_OPTIMIZE_ROUTING_REPLAY : async D→H routing-replay patch (runtime) # - RELAX_OPD_PREEXPANDED_PATCH : OPD pre-expanded multimodal patch (runtime) # - RELAX_OPD_PER_POS_TOKEN_IDS : OPD per-position token_ids logprob; 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..c6504aa91 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,33 @@ 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", "exact"], + default="proxy", + help=( + "P3O behavior-KL implementation. Production training supports proxy and proxy_safe; " + "exact is retained for pure full-vocabulary verification helpers and is rejected by validation." + ), + ) + 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 +2900,97 @@ 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", "exact"}: + raise ValueError(f"--p3o-kl-mode must be proxy, proxy_safe, or exact, got {kl_mode!r}.") + if kl_mode == "exact": + raise ValueError( + "--p3o-kl-mode exact is verifier-only and cannot be used for production P3O training: " + "rollouts store selected-token behavior log-probabilities, not full-vocabulary behavior logits." + ) + 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 +3701,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_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/backends/sglang/test_deterministic_sampler_patch.py b/tests/backends/sglang/test_deterministic_sampler_patch.py new file mode 100644 index 000000000..75ab6add9 --- /dev/null +++ b/tests/backends/sglang/test_deterministic_sampler_patch.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest +import torch + +from relax.backends.sglang import deterministic_sampler_patch as patch + + +def _identity_compile(*, dynamic): + assert dynamic is True + return lambda function: function + + +def _install_fake_sglang_modules( + monkeypatch: pytest.MonkeyPatch, + sampler: ModuleType, + hash_module: ModuleType, +) -> None: + sglang = ModuleType("sglang") + srt = ModuleType("sglang.srt") + layers = ModuleType("sglang.srt.layers") + utils = ModuleType("sglang.srt.layers.utils") + sglang.srt = srt + srt.layers = layers + layers.sampler = sampler + layers.utils = utils + utils.hash = hash_module + monkeypatch.setitem(sys.modules, "sglang", sglang) + monkeypatch.setitem(sys.modules, "sglang.srt", srt) + monkeypatch.setitem(sys.modules, "sglang.srt.layers", layers) + monkeypatch.setitem(sys.modules, "sglang.srt.layers.sampler", sampler) + monkeypatch.setitem(sys.modules, "sglang.srt.layers.utils", utils) + monkeypatch.setitem(sys.modules, "sglang.srt.layers.utils.hash", hash_module) + + +def test_uniform_hash_endpoint_has_finite_upstream_gumbel_cap(): + values = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float64) + + result = patch._uniform_hash_to_gumbel_(values) + + assert torch.isfinite(result).all() + assert result[1].item() == pytest.approx(-torch.log(-torch.log(torch.tensor(0.5))).item()) + assert result[-1].item() == pytest.approx(-torch.log(torch.tensor(2.0**-32)).item()) + + +def test_safe_multinomial_does_not_let_uint32_endpoint_override_logprob(): + uint32_max = torch.iinfo(torch.uint32).max + + def fake_hash(seed, positions, column_indices): + assert seed.shape == positions.shape == (1,) + assert column_indices.shape == (2,) + return torch.tensor([[uint32_max, uint32_max // 2]], dtype=torch.uint32) + + sample = patch._build_safe_multinomial_with_seed(fake_hash, compile_function=_identity_compile) + selected = sample( + torch.tensor([[float("-inf"), 0.0]], dtype=torch.float64), + torch.tensor([44], dtype=torch.int64), + torch.tensor([179], dtype=torch.int64), + ) + + assert selected.tolist() == [[1]] + + +def test_apply_patch_is_version_gated_and_idempotent(monkeypatch): + sampler = ModuleType("sglang.srt.layers.sampler") + + def original(logprobs, seed, positions): + return logprobs, seed, positions + + sampler.multinomial_with_seed = original + hash_module = ModuleType("sglang.srt.layers.utils.hash") + hash_module.murmur_hash32 = object() + _install_fake_sglang_modules(monkeypatch, sampler, hash_module) + monkeypatch.setattr(patch, "_installed_sglang_version", lambda: "0.5.12.post1") + replacement = lambda *_args: None + monkeypatch.setattr(patch, "_build_safe_multinomial_with_seed", lambda _hash: replacement) + + assert patch.apply_deterministic_sampler_endpoint_patch() is True + assert sampler.multinomial_with_seed is replacement + assert getattr(replacement, patch._PATCH_MARKER) is True + assert patch.apply_deterministic_sampler_endpoint_patch() is False + + +def test_apply_patch_leaves_unaffected_sglang_unchanged(monkeypatch): + monkeypatch.setattr(patch, "_installed_sglang_version", lambda: "0.5.13") + + assert patch.apply_deterministic_sampler_endpoint_patch() is False + + +def test_local_version_suffix_resolves_to_affected_public_version(monkeypatch): + monkeypatch.setattr(patch, "version", lambda _package: "0.5.12.post1+cu129") + + assert patch._installed_sglang_version() == "0.5.12.post1" + + +def test_affected_signature_drift_fails_closed(monkeypatch): + sampler = ModuleType("sglang.srt.layers.sampler") + sampler.multinomial_with_seed = lambda inputs, seed: (inputs, seed) + hash_module = ModuleType("sglang.srt.layers.utils.hash") + hash_module.murmur_hash32 = object() + _install_fake_sglang_modules(monkeypatch, sampler, hash_module) + monkeypatch.setattr(patch, "_installed_sglang_version", lambda: "0.5.12.post1") + + with pytest.raises(RuntimeError, match="signature changed"): + patch.apply_deterministic_sampler_endpoint_patch() diff --git a/tests/backends/sglang/test_router_registration.py b/tests/backends/sglang/test_router_registration.py index 84cc8f3e0..201788bc1 100644 --- a/tests/backends/sglang/test_router_registration.py +++ b/tests/backends/sglang/test_router_registration.py @@ -2,6 +2,7 @@ import importlib import logging +import pickle import sys from types import ModuleType, SimpleNamespace @@ -148,6 +149,72 @@ def _make_engine(sglang_engine_module): return engine +@pytest.mark.parametrize("routing_replay", [False, True]) +def test_scheduler_wrapper_applies_endpoint_patch_before_optional_routing_patch( + monkeypatch, sglang_engine_module, routing_replay +): + events = [] + endpoint_patch_module = ModuleType("relax.backends.sglang.deterministic_sampler_patch") + endpoint_patch_module.apply_deterministic_sampler_endpoint_patch = lambda: events.append("endpoint") + routing_patch_module = ModuleType("relax.backends.sglang.routing_replay_patch") + routing_patch_module.apply_patch = lambda: events.append("routing") + scheduler_module = ModuleType("sglang.srt.managers.scheduler") + + def run_scheduler_process(*args, **kwargs): + events.append(("scheduler", args, kwargs)) + return "finished" + + scheduler_module.run_scheduler_process = run_scheduler_process + monkeypatch.setitem(sys.modules, endpoint_patch_module.__name__, endpoint_patch_module) + monkeypatch.setitem(sys.modules, routing_patch_module.__name__, routing_patch_module) + monkeypatch.setitem(sys.modules, scheduler_module.__name__, scheduler_module) + monkeypatch.setattr( + sglang_engine_module.Envs, + "RELAX_OPTIMIZE_ROUTING_REPLAY", + routing_replay, + raising=False, + ) + + assert sglang_engine_module._patched_run_scheduler_process("arg", key="value") == "finished" + expected = ["endpoint"] + if routing_replay: + expected.append("routing") + assert events[:-1] == expected + assert events[-1] == (("scheduler", ("arg",), {"key": "value"})) + + +@pytest.mark.parametrize("routing_replay", [False, True]) +def test_launch_server_always_receives_picklable_scheduler_wrapper(monkeypatch, sglang_engine_module, routing_replay): + calls = [] + http_server = ModuleType("sglang.srt.entrypoints.http_server") + + def launch_server(server_args, **kwargs): + calls.append((server_args, kwargs)) + + http_server.launch_server = launch_server + monkeypatch.setitem(sys.modules, http_server.__name__, http_server) + monkeypatch.setattr(sglang_engine_module.Envs, "RELAX_OPD_PREEXPANDED_PATCH", False, raising=False) + monkeypatch.setattr( + sglang_engine_module.Envs, + "RELAX_OPTIMIZE_ROUTING_REPLAY", + routing_replay, + raising=False, + ) + server_args = object() + + sglang_engine_module._launch_server_with_patches(server_args) + + assert calls == [ + ( + server_args, + {"run_scheduler_process_func": sglang_engine_module._patched_run_scheduler_process}, + ) + ] + assert pickle.loads(pickle.dumps(sglang_engine_module._patched_run_scheduler_process)).__name__ == ( + "_patched_run_scheduler_process" + ) + + def test_missing_load_format_choices_uses_legacy_remote(sglang_engine_module): assert sglang_engine_module._preferred_s3_stream_load_format() == "remote" 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..29f83edbc --- /dev/null +++ b/tests/utils/test_p3o_arguments.py @@ -0,0 +1,187 @@ +# 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_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_exact_kl_before_training(): + with pytest.raises(ValueError, match="verifier-only"): + validate_p3o_args(_p3o_args(p3o_kl_mode="exact")) + + +@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)